mlx_rs/nn/
quantized.rs

1use std::iter::once;
2
3use crate::{
4    array,
5    error::Exception,
6    module::{Module, ModuleParameters, Param},
7    ops::indexing::IndexOp,
8    ops::{self, dequantize, quantized_matmul, zeros},
9    quantization::Quantizable,
10    random::uniform,
11    Array,
12};
13use mlx_internal_macros::{Buildable, Builder};
14use mlx_macros::ModuleParameters;
15
16use crate::nn::{Embedding, Linear};
17
18/// Quantize a module.
19///
20/// # Params
21///
22/// - `module`: The module to quantize.
23/// - `group_size`: The group size to use for the quantized weight. Default to [`Quantizable::DEFAULT_GROUP_SIZE`]
24/// - `bits`: The bit width to use for the quantized weight. Default to [`Quantizable::DEFAULT_BITS`]
25pub fn quantize<M>(
26    module: M,
27    group_size: impl Into<Option<i32>>,
28    bits: impl Into<Option<i32>>,
29) -> Result<M::Quantized, M::QuantizationError>
30where
31    M: Quantizable,
32{
33    let group_size = group_size.into().unwrap_or(M::DEFAULT_GROUP_SIZE);
34    let bits = bits.into().unwrap_or(M::DEFAULT_BITS);
35    module.try_into_quantized(group_size, bits)
36}
37
38/// Builder for [`QuantizedEmbedding`]
39#[derive(Debug, Clone, Builder)]
40#[builder(
41    root = crate,
42    build_with = build_quantized_embedding,
43    err = Exception,
44)]
45pub struct QuantizedEmbeddingBuilder {
46    /// How many possible discrete tokens can we embed. Usually called the vocabulary size.
47    pub embedding_count: i32,
48
49    /// The dimensionality of the embeddings.
50    pub dimensions: i32,
51
52    /// Quantization group size. Default to [`QuantizedEmbedding::DEFAULT_GROUP_SIZE`]
53    #[builder(optional, default = QuantizedEmbedding::DEFAULT_GROUP_SIZE)]
54    pub group_size: i32,
55
56    /// Bits per parameter. Default to [`QuantizedEmbedding::DEFAULT_BITS`]
57    #[builder(optional, default = QuantizedEmbedding::DEFAULT_BITS)]
58    pub bits: i32,
59}
60
61/// The same as ``Embedding`` but with a quantized weight matrix.
62#[derive(Debug, Clone, ModuleParameters, Buildable)]
63#[module(root = crate)]
64#[buildable(root = crate)]
65pub struct QuantizedEmbedding {
66    /// Quantization group size. Default to [`QuantizedEmbedding::DEFAULT_GROUP_SIZE`]
67    pub group_size: i32,
68
69    /// Bits per parameter. Default to [`QuantizedEmbedding::DEFAULT_BITS`]
70    pub bits: i32,
71
72    /// Scales
73    #[param]
74    pub scales: Param<Array>,
75
76    /// Biases
77    #[param]
78    pub biases: Param<Array>,
79
80    /// Inner embedding
81    #[param]
82    pub inner: Embedding,
83}
84
85impl QuantizedEmbeddingBuilder {
86    /// Convenience method to build a new [`QuantizedEmbedding`] with an existing [`Embedding`]
87    pub fn build_with_embedding(
88        self,
89        embedding: Embedding,
90    ) -> Result<QuantizedEmbedding, Exception> {
91        let weight = embedding.weight.value;
92        self.build_with_weight(weight)
93    }
94
95    /// Convenience method to build a new [`QuantizedEmbedding`] with an existing weight matrix
96    pub fn build_with_weight(self, weight: Array) -> Result<QuantizedEmbedding, Exception> {
97        let group_size = self.group_size;
98        let bits = self.bits;
99        build_quantized_embedding_inner(weight, group_size, bits)
100    }
101}
102
103fn build_quantized_embedding_inner(
104    weight: Array,
105    group_size: i32,
106    bits: i32,
107) -> Result<QuantizedEmbedding, Exception> {
108    let (quantized_weight, scales, biases) = ops::quantize(
109        &weight,
110        ops::QuantizeOptions {
111            group_size: Some(group_size),
112            bits: Some(bits),
113            ..Default::default()
114        },
115    )?;
116    let biases = biases.expect("affine quantization returns biases");
117
118    let inner = Embedding {
119        weight: Param::new(quantized_weight),
120    };
121
122    let mut qe = QuantizedEmbedding {
123        group_size,
124        bits,
125        scales: Param::new(scales),
126        biases: Param::new(biases),
127        inner,
128    };
129
130    // Freeze all parameters
131    qe.freeze_parameters(true);
132
133    Ok(qe)
134}
135
136fn build_quantized_embedding(
137    builder: QuantizedEmbeddingBuilder,
138) -> Result<QuantizedEmbedding, Exception> {
139    let embedding_count = builder.embedding_count;
140    let dims = builder.dimensions;
141
142    let scale = array!(f32::sqrt(1.0 / (dims as f32)));
143    // SAFETY: This is safe because the array scale is a single element array
144    let weight = crate::random::normal::<f32>(&[embedding_count, dims], None, None, None)? * &scale;
145
146    builder.build_with_weight(weight)
147}
148
149impl QuantizedEmbedding {
150    /// Default group size
151    pub const DEFAULT_GROUP_SIZE: i32 = 64;
152
153    /// Default bits
154    pub const DEFAULT_BITS: i32 = 4;
155
156    /// Convert an embedding layer to a quantized embedding layer.
157    ///
158    /// # Params
159    ///
160    /// - `embedding`: The embedding layer to convert.
161    /// - `group_size`: The group size to use for the quantized weight. Default to [`QuantizedEmbedding::DEFAULT_GROUP_SIZE`]
162    /// - `bits`: The bit width to use for the quantized weight. Default to [`QuantizedEmbedding::DEFAULT_BITS`]
163    pub fn try_from_embedding(
164        embedding: Embedding,
165        group_size: impl Into<Option<i32>>,
166        bits: impl Into<Option<i32>>,
167    ) -> Result<Self, Exception> {
168        let group_size = group_size.into().unwrap_or(Self::DEFAULT_GROUP_SIZE);
169        let bits = bits.into().unwrap_or(Self::DEFAULT_BITS);
170        build_quantized_embedding_inner(embedding.weight.value, group_size, bits)
171    }
172
173    /// Call the embedding layer as a linear layer.
174    ///
175    /// Use this for example when input embedding and output projection
176    /// weights are tied.
177    pub fn as_linear(&self, x: impl AsRef<Array>) -> Result<Array, Exception> {
178        quantized_matmul(
179            x.as_ref(),
180            &self.inner.weight,
181            &self.scales,
182            self.biases.as_ref(),
183            ops::QuantizedMatmulOptions {
184                transpose: true,
185                group_size: Some(self.group_size),
186                bits: Some(self.bits),
187                ..Default::default()
188            },
189        )
190        .map_err(Into::into)
191    }
192}
193
194impl TryFrom<Embedding> for QuantizedEmbedding {
195    type Error = Exception;
196
197    fn try_from(embedding: Embedding) -> Result<Self, Self::Error> {
198        Self::try_from_embedding(embedding, None, None)
199    }
200}
201
202impl Module<&Array> for QuantizedEmbedding {
203    type Error = Exception;
204    type Output = Array;
205
206    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
207        let s = x.shape();
208        let x = x.flatten(None, None)?;
209        let w = self.inner.weight.index(&x);
210        let scales = self.scales.index(&x);
211        let biases = self.biases.index(&x);
212
213        let out = dequantize(
214            &w,
215            &scales,
216            &biases,
217            ops::DequantizeOptions {
218                group_size: Some(self.group_size),
219                bits: Some(self.bits),
220                ..Default::default()
221            },
222        )?;
223
224        let ret_shape = s.iter().copied().chain(once(-1)).collect::<Vec<_>>();
225        out.reshape(&ret_shape)
226    }
227
228    fn training_mode(&mut self, mode: bool) {
229        self.inner.training_mode(mode);
230    }
231}
232
233/// Builder for [`QuantizedLinear`]
234#[derive(Debug, Clone, Builder)]
235#[builder(
236    root = crate,
237    build_with = build_quantized_linear,
238    err = Exception,
239)]
240pub struct QuantizedLinearBuilder {
241    /// The dimensionality of the input features.
242    pub input_dims: i32,
243
244    /// The dimensionality of the output features.
245    pub output_dims: i32,
246
247    /// Quantization group size. Default to [`QuantizedLinear::DEFAULT_GROUP_SIZE`]
248    #[builder(optional, default = QuantizedLinear::DEFAULT_GROUP_SIZE)]
249    pub group_size: i32,
250
251    /// Bits per parameter. Default to [`QuantizedLinear::DEFAULT_BITS`]
252    #[builder(optional, default = QuantizedLinear::DEFAULT_BITS)]
253    pub bits: i32,
254
255    /// Whether the linear layer has a bias. Default to [`Linear::DEFAULT_BIAS`]
256    #[builder(optional, default = Linear::DEFAULT_BIAS)]
257    pub bias: bool,
258}
259
260impl QuantizedLinearBuilder {
261    /// Convenience method to build a new [`QuantizedLinear`] with an existing [`Linear`]
262    pub fn build_with_linear(self, other: Linear) -> Result<QuantizedLinear, Exception> {
263        self.build_with_weight_and_bias(other.weight.value, other.bias.value)
264    }
265
266    fn build_with_weight_and_bias(
267        self,
268        weight: Array,
269        bias: Option<Array>,
270    ) -> Result<QuantizedLinear, Exception> {
271        build_quantized_linear_inner(weight, bias, self.group_size, self.bits)
272    }
273}
274
275fn build_quantized_linear_inner(
276    weight: Array,
277    bias: Option<Array>,
278    group_size: i32,
279    bits: i32,
280) -> Result<QuantizedLinear, Exception> {
281    let (quantized_weight, scales, biases) = ops::quantize(
282        &weight,
283        ops::QuantizeOptions {
284            group_size: Some(group_size),
285            bits: Some(bits),
286            ..Default::default()
287        },
288    )?;
289    let biases = biases.expect("affine quantization returns biases");
290
291    let inner = Linear {
292        weight: Param::new(quantized_weight),
293        bias: Param::new(bias),
294    };
295
296    let mut ql = QuantizedLinear {
297        group_size,
298        bits,
299        scales: Param::new(scales),
300        biases: Param::new(biases),
301        inner,
302    };
303
304    // Freeze all parameters
305    ql.freeze_parameters(true);
306
307    Ok(ql)
308}
309
310/// Builds a new [`QuantizedLinear`]
311pub fn build_quantized_linear(
312    builder: QuantizedLinearBuilder,
313) -> Result<QuantizedLinear, Exception> {
314    let input_dims = builder.input_dims;
315    let output_dims = builder.output_dims;
316    let scale = f32::sqrt(1.0 / (input_dims as f32));
317    let weight = uniform::<_, f32>(-scale, scale, &[output_dims, input_dims], None)?;
318
319    let bias = if builder.bias {
320        Some(zeros::<f32>(&[output_dims])?)
321    } else {
322        None
323    };
324
325    builder.build_with_weight_and_bias(weight, bias)
326}
327
328/// Applies an affine transformation to the input using a quantized weight matrix.
329///
330/// It is the quantized equivalent of [`Linear`].  For now its
331/// parameters are frozen and will not be included in any gradient computation
332/// but this will probably change in the future.
333///
334/// QuantizedLinear also provides several useful static to convert linear
335/// layers to QuantizedLinear layers.
336#[derive(Debug, Clone, ModuleParameters, Buildable)]
337#[module(root = crate)]
338#[buildable(root = crate)]
339pub struct QuantizedLinear {
340    /// Quantization group size. Default to [`QuantizedLinear::DEFAULT_GROUP_SIZE`]
341    pub group_size: i32,
342
343    /// Bits per parameter. Default to [`QuantizedLinear::DEFAULT_BITS`]
344    pub bits: i32,
345
346    /// Scales
347    #[param]
348    pub scales: Param<Array>,
349
350    /// Biases
351    #[param]
352    pub biases: Param<Array>,
353
354    /// Inner linear layer
355    #[param]
356    pub inner: Linear,
357}
358
359impl QuantizedLinear {
360    /// Default group size
361    pub const DEFAULT_GROUP_SIZE: i32 = 64;
362
363    /// Default bits
364    pub const DEFAULT_BITS: i32 = 4;
365
366    /// Convert a linear layer to a quantized linear layer.
367    ///
368    /// # Params
369    ///
370    /// - `linear`: The linear layer to convert.
371    /// - `group_size`: The group size to use for the quantized weight. Default to [`QuantizedLinear::DEFAULT_GROUP_SIZE`]
372    /// - `bits`: The bit width to use for the quantized weight. Default to [`QuantizedLinear::DEFAULT_BITS`]
373    pub fn try_from_linear(
374        linear: Linear,
375        group_size: impl Into<Option<i32>>,
376        bits: impl Into<Option<i32>>,
377    ) -> Result<Self, Exception> {
378        let group_size = group_size.into().unwrap_or(Self::DEFAULT_GROUP_SIZE);
379        let bits = bits.into().unwrap_or(Self::DEFAULT_BITS);
380        build_quantized_linear_inner(linear.weight.value, linear.bias.value, group_size, bits)
381    }
382}
383
384impl TryFrom<Linear> for QuantizedLinear {
385    type Error = Exception;
386
387    fn try_from(linear: Linear) -> Result<Self, Self::Error> {
388        Self::try_from_linear(linear, None, None)
389    }
390}
391
392impl Module<&Array> for QuantizedLinear {
393    type Error = Exception;
394    type Output = Array;
395
396    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
397        let mut x = quantized_matmul(
398            x,
399            &self.inner.weight,
400            &self.scales,
401            self.biases.as_ref(),
402            ops::QuantizedMatmulOptions {
403                transpose: true,
404                group_size: Some(self.group_size),
405                bits: Some(self.bits),
406                ..Default::default()
407            },
408        )?;
409        if let Some(bias) = &self.inner.bias.value {
410            x = x.add(bias)?;
411        }
412        Ok(x)
413    }
414
415    fn training_mode(&mut self, mode: bool) {
416        self.inner.training_mode(mode);
417    }
418}
mlx_rs/nn/quantized.rs

mlx_rs/nn/
quantized.rs
