candle_transformers/models/
quantized_llama2_c.rs

1//! Quantized Llama2 model implementation.
2//!
3//! This provides an 8-bit quantized implementation of Meta's LLaMA2 language model
4//! for reduced memory usage and faster inference.
5//!
6//! Key characteristics:
7//! - Decoder-only transformer architecture
8//! - RoPE position embeddings
9//! - Grouped Query Attention
10//! - 8-bit quantization of weights
11//!
12//! References:
13//! - [LLaMA2 Paper](https://arxiv.org/abs/2307.09288)
14//! - [LLaMA2 Technical Report](https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/)
15//!
16
17use super::llama2_c::{Cache, Config};
18use crate::quantized_nn::{linear_no_bias as linear, Embedding, Linear, RmsNorm};
19pub use crate::quantized_var_builder::VarBuilder;
20use candle::{DType, IndexOp, Module, Result, Tensor, D};
21
22fn silu(xs: &Tensor) -> Result<Tensor> {
23    xs / (xs.neg()?.exp()? + 1.0)?
24}
25
26#[derive(Debug, Clone)]
27struct CausalSelfAttention {
28    q_proj: Linear,
29    k_proj: Linear,
30    v_proj: Linear,
31    o_proj: Linear,
32    n_head: usize,
33    n_key_value_head: usize,
34    head_dim: usize,
35}
36
37impl CausalSelfAttention {
38    fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result<Tensor> {
39        let (b_sz, seq_len, h, n_embd) = x.dims4()?;
40        let cos = cache.cos.i(index_pos..index_pos + seq_len)?;
41        let sin = cache.sin.i(index_pos..index_pos + seq_len)?;
42        let cos = cos.unsqueeze(1)?;
43        let sin = sin.unsqueeze(1)?;
44        let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
45        let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
46        let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 2))?;
47        let x0 = x.narrow(D::Minus1, 0, 1)?;
48        let x1 = x.narrow(D::Minus1, 1, 1)?;
49        let dst0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?;
50        let dst1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?;
51        let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?.reshape((b_sz, seq_len, h, n_embd))?;
52        Ok(rope)
53    }
54
55    fn forward(
56        &self,
57        x: &Tensor,
58        index_pos: usize,
59        block_idx: usize,
60        cache: &mut Cache,
61    ) -> Result<Tensor> {
62        let (b_sz, seq_len, n_embd) = x.dims3()?;
63        let q = self.q_proj.forward(x)?;
64        let k = self.k_proj.forward(x)?;
65        let v = self.v_proj.forward(x)?;
66
67        let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?;
68        let k = k.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?;
69        let mut v = v.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?;
70
71        let q = self.apply_rotary_emb(&q, index_pos, cache)?;
72        let mut k = self.apply_rotary_emb(&k, index_pos, cache)?;
73
74        if cache.use_kv_cache {
75            if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] {
76                k = Tensor::cat(&[cache_k, &k], 1)?.contiguous()?;
77                v = Tensor::cat(&[cache_v, &v], 1)?.contiguous()?;
78            }
79            cache.kvs[block_idx] = Some((k.clone(), v.clone()))
80        }
81
82        let k = self.repeat_kv(k)?;
83        let v = self.repeat_kv(v)?;
84
85        let q = q.transpose(1, 2)?.contiguous()?;
86        let k = k.transpose(1, 2)?.contiguous()?;
87        let v = v.transpose(1, 2)?.contiguous()?;
88
89        let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
90        let att = if seq_len <= 1 {
91            att
92        } else {
93            let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
94            masked_fill(&att, &mask, f32::NEG_INFINITY)?
95        };
96        let att = candle_nn::ops::softmax(&att, D::Minus1)?;
97        // Convert to contiguous as matmul doesn't support strided vs for now.
98        let y = att.matmul(&v.contiguous()?)?;
99        let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
100        let y = self.o_proj.forward(&y)?;
101        Ok(y)
102    }
103
104    fn repeat_kv(&self, x: Tensor) -> Result<Tensor> {
105        let n_rep = self.n_head / self.n_key_value_head;
106        if n_rep == 1 {
107            Ok(x)
108        } else {
109            let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?;
110            let x = x
111                .unsqueeze(3)?
112                .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))?
113                .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?;
114            Ok(x)
115        }
116    }
117
118    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
119        let size_in = cfg.dim;
120        let size_q = (cfg.dim / cfg.n_heads) * cfg.n_heads;
121        let size_kv = (cfg.dim / cfg.n_heads) * cfg.n_kv_heads;
122        let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?;
123        let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?;
124        let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?;
125        let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?;
126        Ok(Self {
127            q_proj,
128            k_proj,
129            v_proj,
130            o_proj,
131            n_head: cfg.n_heads,
132            n_key_value_head: cfg.n_kv_heads,
133            head_dim: cfg.dim / cfg.n_heads,
134        })
135    }
136}
137
138fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
139    let shape = mask.shape();
140    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
141    let m = mask.where_cond(&on_true, on_false)?;
142    Ok(m)
143}
144
145#[derive(Debug, Clone)]
146struct Mlp {
147    c_fc1: Linear,
148    c_fc2: Linear,
149    c_proj: Linear,
150}
151
152impl Mlp {
153    fn new(c_fc1: Linear, c_fc2: Linear, c_proj: Linear) -> Self {
154        Self {
155            c_fc1,
156            c_fc2,
157            c_proj,
158        }
159    }
160
161    fn forward(&self, x: &Tensor) -> Result<Tensor> {
162        let x = (silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?;
163        self.c_proj.forward(&x)
164    }
165
166    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
167        let h_size = cfg.dim;
168        let i_size = cfg.hidden_dim;
169        let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?;
170        let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?;
171        let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?;
172        Ok(Self::new(c_fc1, c_fc2, c_proj))
173    }
174}
175
176#[derive(Debug, Clone)]
177struct Block {
178    rms_1: RmsNorm,
179    attn: CausalSelfAttention,
180    rms_2: RmsNorm,
181    mlp: Mlp,
182}
183
184impl Block {
185    fn new(rms_1: RmsNorm, attn: CausalSelfAttention, rms_2: RmsNorm, mlp: Mlp) -> Self {
186        Self {
187            rms_1,
188            attn,
189            rms_2,
190            mlp,
191        }
192    }
193
194    fn forward(
195        &self,
196        x: &Tensor,
197        index_pos: usize,
198        block_idx: usize,
199        cache: &mut Cache,
200    ) -> Result<Tensor> {
201        let residual = x;
202        let x = self.rms_1.forward(x)?;
203        let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?;
204        let residual = &x;
205        let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?;
206        Ok(x)
207    }
208
209    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
210        let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?;
211        let mlp = Mlp::load(vb.pp("mlp"), cfg)?;
212        let input_layernorm = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("input_layernorm"))?;
213        let post_attention_layernorm =
214            RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("post_attention_layernorm"))?;
215        Ok(Self::new(
216            input_layernorm,
217            attn,
218            post_attention_layernorm,
219            mlp,
220        ))
221    }
222}
223
224#[derive(Debug, Clone)]
225pub struct QLlama {
226    wte: Embedding,
227    blocks: Vec<Block>,
228    ln_f: RmsNorm,
229    lm_head: Linear,
230    pub config: Config,
231}
232
233impl QLlama {
234    pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result<Tensor> {
235        let (_b_sz, _seq_len) = x.dims2()?;
236        let mut x = self.wte.forward(x)?;
237        for (block_idx, block) in self.blocks.iter().enumerate() {
238            x = block.forward(&x, index_pos, block_idx, cache)?;
239        }
240        let x = self.ln_f.forward(&x)?;
241        let logits = self.lm_head.forward(&x)?;
242        logits.to_dtype(DType::F32)
243    }
244
245    pub fn load(vb: VarBuilder, cfg: Config) -> Result<Self> {
246        let wte = Embedding::new(cfg.vocab_size, cfg.dim, vb.pp("model.embed_tokens"))?;
247        let lm_head = linear(cfg.dim, cfg.vocab_size, vb.pp("lm_head"))?;
248        let ln_f = RmsNorm::new(cfg.dim, cfg.norm_eps, vb.pp("model.norm"))?;
249        let blocks: Vec<_> = (0..cfg.n_layers)
250            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), &cfg).unwrap())
251            .collect();
252        Ok(Self {
253            wte,
254            blocks,
255            ln_f,
256            lm_head,
257            config: cfg,
258        })
259    }
260}
candle_transformers/models/quantized_llama2_c.rs

candle_transformers/models/
quantized_llama2_c.rs
