candle_transformers/models/
quantized_phi3.rs

1//! Phi3 model implementation with quantization support.
2//!
3//! Phi3 is a language model intended for research purposes.
4//! This implementation provides quantization for reduced memory usage.
5//!
6//! Key characteristics:
7//! - Multi-head attention
8//! - RMSNorm for layer normalization
9//! - Rotary positional embeddings (RoPE)
10//! - Support for quantization
11//!
12//! References:
13//! - [Model Card](https://huggingface.co/microsoft/phi-3)
14//!
15
16use std::collections::HashMap;
17
18use candle::quantized::gguf_file;
19use candle::quantized::QTensor;
20use candle::{DType, Device, IndexOp, Module, Result, Tensor, D};
21use candle_nn::{kv_cache::KvCache, Embedding, RmsNorm};
22
23#[derive(Debug, Clone)]
24struct QLinear {
25    inner: candle::quantized::QMatMul,
26    span: tracing::Span,
27}
28
29impl QLinear {
30    fn new<R: std::io::Read + std::io::Seek>(
31        ct: &gguf_file::Content,
32        r: &mut R,
33        name: &str,
34        device: &Device,
35    ) -> Result<Self> {
36        let span = tracing::span!(tracing::Level::TRACE, "qmatmul");
37        let w = ct.tensor(r, &format!("{name}.weight"), device)?;
38        let inner = candle::quantized::QMatMul::from_qtensor(w)?;
39        Ok(Self { inner, span })
40    }
41}
42
43impl Module for QLinear {
44    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
45        let _enter = self.span.enter();
46        self.inner.forward(xs)
47    }
48}
49
50#[derive(Debug, Clone)]
51struct Mlp {
52    ffn_up: QLinear,
53    ffn_down: QLinear,
54    i_size: usize,
55}
56
57impl Module for Mlp {
58    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
59        let up_states = xs.apply(&self.ffn_up)?;
60        let gate = up_states.narrow(D::Minus1, 0, self.i_size)?;
61        let up_states = up_states.narrow(D::Minus1, self.i_size, self.i_size)?;
62        let up_states = (up_states * gate.silu()?)?;
63        up_states.apply(&self.ffn_down)
64    }
65}
66
67fn rms_norm(w: QTensor, eps: f64) -> Result<RmsNorm> {
68    let w = w.dequantize(&w.device())?;
69    let rms = RmsNorm::new(w, eps);
70    Ok(rms)
71}
72
73#[derive(Debug, Clone)]
74struct LayerWeights {
75    attn_qkv: QLinear,
76    attn_output: QLinear,
77    attn_norm: RmsNorm,
78    ffn_norm: RmsNorm,
79    mlp: Mlp,
80    n_head: usize,
81    n_kv_head: usize,
82    head_dim: usize,
83    cos: Tensor,
84    sin: Tensor,
85    neg_inf: Tensor,
86    kv_cache: KvCache,
87    use_flash_attn: bool,
88    span_attn: tracing::Span,
89    span_rot: tracing::Span,
90}
91
92fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result<Tensor> {
93    let shape = mask.shape();
94    let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?;
95    Ok(m)
96}
97
98impl LayerWeights {
99    fn apply_rotary_emb(&self, xs: &Tensor, index_pos: usize) -> Result<Tensor> {
100        let _enter = self.span_rot.enter();
101        let (_b_sz, _h, seq_len, _n_embd) = xs.dims4()?;
102        let cos = self.cos.narrow(0, index_pos, seq_len)?;
103        let sin = self.sin.narrow(0, index_pos, seq_len)?;
104        candle_nn::rotary_emb::rope(&xs.contiguous()?, &cos, &sin)
105    }
106
107    fn forward_attn(
108        &mut self,
109        x: &Tensor,
110        mask: Option<&Tensor>,
111        index_pos: usize,
112    ) -> Result<Tensor> {
113        let _enter = self.span_attn.enter();
114        let (b_sz, seq_len, n_embd) = x.dims3()?;
115        let qkv = self.attn_qkv.forward(x)?;
116
117        let query_pos = self.n_head * self.head_dim;
118        let q = qkv.narrow(D::Minus1, 0, query_pos)?;
119        let k = qkv.narrow(D::Minus1, query_pos, self.n_kv_head * self.head_dim)?;
120        let v = qkv.narrow(
121            D::Minus1,
122            query_pos + self.n_kv_head * self.head_dim,
123            self.n_kv_head * self.head_dim,
124        )?;
125
126        let q = q
127            .reshape((b_sz, seq_len, self.n_head, self.head_dim))?
128            .transpose(1, 2)?;
129        let k = k
130            .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?
131            .transpose(1, 2)?;
132        let v = v
133            .reshape((b_sz, seq_len, self.n_kv_head, self.head_dim))?
134            .transpose(1, 2)?;
135
136        let q = self.apply_rotary_emb(&q, index_pos)?.contiguous()?;
137        let k = self.apply_rotary_emb(&k, index_pos)?;
138
139        let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?;
140
141        let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?;
142        let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?;
143
144        let y = if self.use_flash_attn {
145            // flash-attn expects (b_sz, seq_len, nheads, head_dim)
146            let q = q.to_dtype(DType::BF16)?.transpose(1, 2)?;
147            let k = k.to_dtype(DType::BF16)?.transpose(1, 2)?;
148            let v = v.to_dtype(DType::BF16)?.transpose(1, 2)?;
149            let softmax_scale = 1f32 / (self.head_dim as f32).sqrt();
150            flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?
151                .to_dtype(DType::F32)?
152                .transpose(1, 2)?
153        } else {
154            let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
155            let att = match mask {
156                None => att,
157                Some(mask) => {
158                    let mask = mask.broadcast_as(att.shape())?;
159                    masked_fill(&att, &mask, &self.neg_inf)?
160                }
161            };
162            let att = candle_nn::ops::softmax_last_dim(&att)?;
163            // Convert to contiguous as matmul doesn't support strided vs for now.
164            att.matmul(&v)?
165        };
166        let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
167        let y = self.attn_output.forward(&y)?;
168        Ok(y)
169    }
170}
171
172#[cfg(feature = "flash-attn")]
173fn flash_attn(
174    q: &Tensor,
175    k: &Tensor,
176    v: &Tensor,
177    softmax_scale: f32,
178    causal: bool,
179) -> Result<Tensor> {
180    candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal)
181}
182
183#[cfg(not(feature = "flash-attn"))]
184fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result<Tensor> {
185    unimplemented!("compile with '--features flash-attn'")
186}
187
188#[derive(Debug, Clone)]
189pub struct ModelWeights {
190    tok_embeddings: Embedding,
191    layers: Vec<LayerWeights>,
192    output_norm: RmsNorm,
193    output: QLinear,
194    masks: HashMap<usize, Tensor>,
195    span: tracing::Span,
196    span_output: tracing::Span,
197}
198
199fn precomput_freqs_cis(
200    head_dim: usize,
201    max_seq_len: usize,
202    freq_base: f32,
203    device: &Device,
204) -> Result<(Tensor, Tensor)> {
205    let theta: Vec<_> = (0..head_dim)
206        .step_by(2)
207        .map(|i| 1f32 / freq_base.powf(i as f32 / head_dim as f32))
208        .collect();
209    let theta = Tensor::new(theta.as_slice(), device)?;
210    let idx_theta = Tensor::arange(0, max_seq_len as u32, device)?
211        .to_dtype(DType::F32)?
212        .reshape((max_seq_len, 1))?
213        .matmul(&theta.reshape((1, theta.elem_count()))?)?;
214    let cos = idx_theta.cos()?;
215    let sin = idx_theta.sin()?;
216    Ok((cos, sin))
217}
218
219impl ModelWeights {
220    pub fn from_gguf<R: std::io::Seek + std::io::Read>(
221        use_flash_attn: bool,
222        ct: gguf_file::Content,
223        reader: &mut R,
224        device: &Device,
225    ) -> Result<Self> {
226        let md_get = |s: &str| match ct.metadata.get(s) {
227            None => candle::bail!("cannot find {s} in metadata"),
228            Some(v) => Ok(v),
229        };
230
231        // Parameter extraction from metadata.
232        let head_count = md_get("phi3.attention.head_count")?.to_u32()? as usize;
233        let head_count_kv = md_get("phi3.attention.head_count_kv")?.to_u32()? as usize;
234        let block_count = md_get("phi3.block_count")?.to_u32()? as usize;
235        let embedding_length = md_get("phi3.embedding_length")?.to_u32()? as usize;
236        let max_seq_len = md_get("phi3.context_length")?.to_u32()? as usize;
237        let head_dim = embedding_length / head_count;
238        let i_size = md_get("phi3.feed_forward_length")?.to_u32()? as usize;
239        let rope_dim = md_get("phi3.rope.dimension_count")?.to_u32()? as usize;
240        let rms_eps = md_get("phi3.attention.layer_norm_rms_epsilon")?.to_f32()? as f64;
241        let (cos, sin) = precomput_freqs_cis(rope_dim, max_seq_len, 10_000., device)?;
242        let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?;
243
244        let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?;
245        let tok_embeddings = tok_embeddings.dequantize(device)?;
246        let output_norm = rms_norm(ct.tensor(reader, "output_norm.weight", device)?, rms_eps)?;
247        let output = QLinear::new(&ct, reader, "output", device)?;
248
249        let mut layers = Vec::with_capacity(block_count);
250        for layer_idx in 0..block_count {
251            let prefix = format!("blk.{layer_idx}");
252            let ffn_up = QLinear::new(&ct, reader, &format!("{prefix}.ffn_up"), device)?;
253            let ffn_down = QLinear::new(&ct, reader, &format!("{prefix}.ffn_down"), device)?;
254            let mlp = Mlp {
255                ffn_up,
256                ffn_down,
257                i_size,
258            };
259            let attn_norm = rms_norm(
260                ct.tensor(reader, &format!("{prefix}.attn_norm.weight"), device)?,
261                rms_eps,
262            )?;
263            let ffn_norm = rms_norm(
264                ct.tensor(reader, &format!("{prefix}.ffn_norm.weight"), device)?,
265                rms_eps,
266            )?;
267            let span_attn = tracing::span!(tracing::Level::TRACE, "attn");
268            let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot");
269            let kv_cache = KvCache::new(2, max_seq_len);
270            layers.push(LayerWeights {
271                attn_qkv: QLinear::new(&ct, reader, &format!("{prefix}.attn_qkv"), device)?,
272                attn_output: QLinear::new(&ct, reader, &format!("{prefix}.attn_output"), device)?,
273                attn_norm,
274                ffn_norm,
275                mlp,
276                n_head: head_count,
277                n_kv_head: head_count_kv,
278                head_dim,
279                cos: cos.clone(),
280                sin: sin.clone(),
281                neg_inf: neg_inf.clone(),
282                kv_cache,
283                use_flash_attn,
284                span_attn,
285                span_rot,
286            })
287        }
288        let span = tracing::span!(tracing::Level::TRACE, "model");
289        let span_output = tracing::span!(tracing::Level::TRACE, "output");
290        Ok(Self {
291            tok_embeddings: Embedding::new(tok_embeddings, embedding_length),
292            layers,
293            output_norm,
294            output,
295            masks: HashMap::new(),
296            span,
297            span_output,
298        })
299    }
300
301    fn mask(&mut self, t: usize, device: &Device) -> Result<Tensor> {
302        if let Some(mask) = self.masks.get(&t) {
303            Ok(mask.clone())
304        } else {
305            let mask: Vec<_> = (0..t)
306                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
307                .collect();
308            let mask = Tensor::from_slice(&mask, (t, t), device)?;
309            self.masks.insert(t, mask.clone());
310            Ok(mask)
311        }
312    }
313
314    pub fn forward(&mut self, xs: &Tensor, index_pos: usize) -> Result<Tensor> {
315        let (_b_sz, seq_len) = xs.dims2()?;
316        let mask = if seq_len == 1 {
317            None
318        } else {
319            Some(self.mask(seq_len, xs.device())?)
320        };
321        let _enter = self.span.enter();
322        let mut xs = self.tok_embeddings.forward(xs)?;
323        for layer in self.layers.iter_mut() {
324            let residual = &xs;
325            let ys = xs.apply(&layer.attn_norm)?;
326            let ys = layer.forward_attn(&ys, mask.as_ref(), index_pos)?;
327            let ys = (ys + residual)?;
328            let residual = &ys;
329            let ys = ys.apply(&layer.ffn_norm)?;
330            let ys = layer.mlp.forward(&ys)?;
331            xs = (ys + residual)?
332        }
333        let xs = xs.apply(&self.output_norm)?.i((.., seq_len - 1, ..))?;
334        let _enter = self.span_output.enter();
335        self.output.forward(&xs)
336    }
337}
candle_transformers/models/quantized_phi3.rs

candle_transformers/models/
quantized_phi3.rs
