candle_transformers/models/
llama2_c.rs

1//! Llama2 inference implementation.
2//!
3//! See ["LLaMA 2: Open Foundation and Fine-Tuned Chat Models"](https://arxiv.org/abs/2307.09288)
4//!
5//! - ⚡ [Interactive Wasm Example](https://huggingface.co/spaces/lmz/candle-llama2)
6//! - 💻 llama2.c [GH Link](https://github.com/karpathy/llama2.c)
7//!
8
9use candle::{DType, Device, IndexOp, Result, Tensor, D};
10use candle_nn::linear_no_bias as linear;
11use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder};
12use std::collections::HashMap;
13
14#[derive(Debug, Clone)]
15pub struct Config {
16    pub dim: usize,        // transformer dimension
17    pub hidden_dim: usize, // for ffn layers
18    pub n_layers: usize,   // number of layers
19    pub n_heads: usize,    // number of query heads
20    pub n_kv_heads: usize, // number of key/value heads (can be < query heads because of multiquery)
21    pub vocab_size: usize, // vocabulary size, usually 256 (byte-level)
22    pub seq_len: usize,    // max sequence length
23    pub norm_eps: f64,
24}
25
26impl Config {
27    pub fn tiny_260k() -> Self {
28        Self {
29            dim: 64,
30            hidden_dim: 768,
31            n_layers: 5,
32            n_heads: 8,
33            n_kv_heads: 4,
34            vocab_size: 32000,
35            seq_len: 512,
36            norm_eps: 1e-5,
37        }
38    }
39
40    pub fn tiny_15m() -> Self {
41        Self {
42            dim: 288,
43            hidden_dim: 768,
44            n_layers: 6,
45            n_heads: 6,
46            n_kv_heads: 6,
47            vocab_size: 32000,
48            seq_len: 256,
49            norm_eps: 1e-5,
50        }
51    }
52
53    pub fn tiny_42m() -> Self {
54        Self {
55            dim: 512,
56            hidden_dim: 768,
57            n_layers: 8,
58            n_heads: 8,
59            n_kv_heads: 8,
60            vocab_size: 32000,
61            seq_len: 1024,
62            norm_eps: 1e-5,
63        }
64    }
65
66    pub fn tiny_110m() -> Self {
67        Self {
68            dim: 768,
69            hidden_dim: 768,
70            n_layers: 12,
71            n_heads: 12,
72            n_kv_heads: 12,
73            vocab_size: 32000,
74            seq_len: 1024,
75            norm_eps: 1e-5,
76        }
77    }
78}
79
80#[derive(Debug, Clone)]
81pub struct Cache {
82    masks: HashMap<usize, Tensor>,
83    pub use_kv_cache: bool,
84    pub kvs: Vec<Option<(Tensor, Tensor)>>,
85    pub cos: Tensor,
86    pub sin: Tensor,
87    device: Device,
88}
89
90impl Cache {
91    pub fn new(use_kv_cache: bool, cfg: &Config, vb: VarBuilder) -> Result<Self> {
92        let n_elem = cfg.dim / cfg.n_heads;
93        let theta: Vec<_> = (0..n_elem)
94            .step_by(2)
95            .map(|i| 1f32 / 10000f32.powf(i as f32 / n_elem as f32))
96            .collect();
97        let theta = Tensor::new(theta.as_slice(), vb.device())?;
98        let idx_theta = Tensor::arange(0, cfg.seq_len as u32, vb.device())?
99            .to_dtype(DType::F32)?
100            .reshape((cfg.seq_len, 1))?
101            .matmul(&theta.reshape((1, theta.elem_count()))?)?;
102        let precomputed_cos = idx_theta.cos()?;
103        let precomputed_sin = idx_theta.sin()?;
104
105        let freq_cis_real = vb
106            .get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_real")
107            .unwrap_or(precomputed_cos);
108        let freq_cis_imag = vb
109            .get((cfg.seq_len, cfg.head_size() / 2), "freq_cis_imag")
110            .unwrap_or(precomputed_sin);
111        let cos = freq_cis_real.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?;
112        let sin = freq_cis_imag.reshape((cfg.seq_len, cfg.head_size() / 2, 1))?;
113        Ok(Self {
114            masks: HashMap::new(),
115            use_kv_cache,
116            kvs: vec![None; cfg.n_layers],
117            cos,
118            sin,
119            device: vb.device().clone(),
120        })
121    }
122
123    pub fn mask(&mut self, t: usize) -> Result<Tensor> {
124        if let Some(mask) = self.masks.get(&t) {
125            Ok(mask.clone())
126        } else {
127            let mask: Vec<_> = (0..t)
128                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
129                .collect();
130            let mask = Tensor::from_slice(&mask, (t, t), &self.device)?;
131            self.masks.insert(t, mask.clone());
132            Ok(mask)
133        }
134    }
135}
136
137fn silu(xs: &Tensor) -> Result<Tensor> {
138    xs / (xs.neg()?.exp()? + 1.0)?
139}
140
141#[derive(Debug, Clone)]
142struct CausalSelfAttention {
143    q_proj: Linear,
144    k_proj: Linear,
145    v_proj: Linear,
146    o_proj: Linear,
147    n_head: usize,
148    n_key_value_head: usize,
149    head_dim: usize,
150}
151
152impl CausalSelfAttention {
153    fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result<Tensor> {
154        let (b_sz, seq_len, h, n_embd) = x.dims4()?;
155        let cos = cache.cos.i(index_pos..index_pos + seq_len)?;
156        let sin = cache.sin.i(index_pos..index_pos + seq_len)?;
157        let cos = cos.unsqueeze(1)?;
158        let sin = sin.unsqueeze(1)?;
159        let cos = cos.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
160        let sin = sin.broadcast_as((b_sz, seq_len, 1, n_embd / 2, 1))?;
161        let x = x.reshape((b_sz, seq_len, h, n_embd / 2, 2))?;
162        let x0 = x.narrow(D::Minus1, 0, 1)?;
163        let x1 = x.narrow(D::Minus1, 1, 1)?;
164        let dst0 = (x0.broadcast_mul(&cos)? - x1.broadcast_mul(&sin)?)?;
165        let dst1 = (x0.broadcast_mul(&sin)? + x1.broadcast_mul(&cos)?)?;
166        let rope = Tensor::cat(&[&dst0, &dst1], D::Minus1)?.reshape((b_sz, seq_len, h, n_embd))?;
167        Ok(rope)
168    }
169
170    fn forward(
171        &self,
172        x: &Tensor,
173        index_pos: usize,
174        block_idx: usize,
175        cache: &mut Cache,
176    ) -> Result<Tensor> {
177        let (b_sz, seq_len, n_embd) = x.dims3()?;
178        let q = self.q_proj.forward(x)?;
179        let k = self.k_proj.forward(x)?;
180        let v = self.v_proj.forward(x)?;
181
182        let q = q.reshape((b_sz, seq_len, self.n_head, self.head_dim))?;
183        let k = k.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?;
184        let mut v = v.reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))?;
185
186        let q = self.apply_rotary_emb(&q, index_pos, cache)?;
187        let mut k = self.apply_rotary_emb(&k, index_pos, cache)?;
188
189        if cache.use_kv_cache {
190            if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] {
191                k = Tensor::cat(&[cache_k, &k], 1)?.contiguous()?;
192                v = Tensor::cat(&[cache_v, &v], 1)?.contiguous()?;
193            }
194            cache.kvs[block_idx] = Some((k.clone(), v.clone()))
195        }
196
197        let k = self.repeat_kv(k)?;
198        let v = self.repeat_kv(v)?;
199
200        let q = q.transpose(1, 2)?.contiguous()?;
201        let k = k.transpose(1, 2)?.contiguous()?;
202        let v = v.transpose(1, 2)?.contiguous()?;
203
204        let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
205        let att = if seq_len <= 1 {
206            att
207        } else {
208            let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
209            masked_fill(&att, &mask, f32::NEG_INFINITY)?
210        };
211        let att = candle_nn::ops::softmax(&att, D::Minus1)?;
212        // Convert to contiguous as matmul doesn't support strided vs for now.
213        let y = att.matmul(&v.contiguous()?)?;
214        let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?;
215        let y = self.o_proj.forward(&y)?;
216        Ok(y)
217    }
218
219    fn repeat_kv(&self, x: Tensor) -> Result<Tensor> {
220        let n_rep = self.n_head / self.n_key_value_head;
221        if n_rep == 1 {
222            Ok(x)
223        } else {
224            let (b_sz, seq_len, n_kv_head, head_dim) = x.dims4()?;
225            let x = x
226                .unsqueeze(3)?
227                .expand((b_sz, seq_len, n_kv_head, n_rep, head_dim))?
228                .reshape((b_sz, seq_len, n_kv_head * n_rep, head_dim))?;
229            Ok(x)
230        }
231    }
232
233    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
234        let size_in = cfg.dim;
235        let size_q = (cfg.dim / cfg.n_heads) * cfg.n_heads;
236        let size_kv = (cfg.dim / cfg.n_heads) * cfg.n_kv_heads;
237        let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?;
238        let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?;
239        let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?;
240        let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?;
241        Ok(Self {
242            q_proj,
243            k_proj,
244            v_proj,
245            o_proj,
246            n_head: cfg.n_heads,
247            n_key_value_head: cfg.n_kv_heads,
248            head_dim: cfg.dim / cfg.n_heads,
249        })
250    }
251}
252
253fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
254    let shape = mask.shape();
255    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
256    let m = mask.where_cond(&on_true, on_false)?;
257    Ok(m)
258}
259
260#[derive(Debug, Clone)]
261struct Mlp {
262    c_fc1: Linear,
263    c_fc2: Linear,
264    c_proj: Linear,
265}
266
267impl Mlp {
268    fn new(c_fc1: Linear, c_fc2: Linear, c_proj: Linear) -> Self {
269        Self {
270            c_fc1,
271            c_fc2,
272            c_proj,
273        }
274    }
275
276    fn forward(&self, x: &Tensor) -> Result<Tensor> {
277        let x = (silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?;
278        self.c_proj.forward(&x)
279    }
280
281    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
282        let h_size = cfg.dim;
283        let i_size = cfg.hidden_dim;
284        let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?;
285        let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?;
286        let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?;
287        Ok(Self::new(c_fc1, c_fc2, c_proj))
288    }
289}
290
291#[derive(Debug, Clone)]
292struct Block {
293    rms_1: RmsNorm,
294    attn: CausalSelfAttention,
295    rms_2: RmsNorm,
296    mlp: Mlp,
297}
298
299impl Block {
300    fn new(rms_1: RmsNorm, attn: CausalSelfAttention, rms_2: RmsNorm, mlp: Mlp) -> Self {
301        Self {
302            rms_1,
303            attn,
304            rms_2,
305            mlp,
306        }
307    }
308
309    fn forward(
310        &self,
311        x: &Tensor,
312        index_pos: usize,
313        block_idx: usize,
314        cache: &mut Cache,
315    ) -> Result<Tensor> {
316        let residual = x;
317        let x = self.rms_1.forward(x)?;
318        let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?;
319        let residual = &x;
320        let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?;
321        Ok(x)
322    }
323
324    fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
325        let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?;
326        let mlp = Mlp::load(vb.pp("mlp"), cfg)?;
327        let input_layernorm = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("input_layernorm"))?;
328        let post_attention_layernorm =
329            rms_norm(cfg.dim, cfg.norm_eps, vb.pp("post_attention_layernorm"))?;
330        Ok(Self::new(
331            input_layernorm,
332            attn,
333            post_attention_layernorm,
334            mlp,
335        ))
336    }
337}
338
339#[derive(Debug, Clone)]
340pub struct Llama {
341    wte: Embedding,
342    blocks: Vec<Block>,
343    ln_f: RmsNorm,
344    lm_head: Linear,
345    pub config: Config,
346}
347
348impl Llama {
349    pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result<Tensor> {
350        let (_b_sz, _seq_len) = x.dims2()?;
351        let mut x = self.wte.forward(x)?;
352        for (block_idx, block) in self.blocks.iter().enumerate() {
353            x = block.forward(&x, index_pos, block_idx, cache)?;
354        }
355        let x = self.ln_f.forward(&x)?;
356        let logits = self.lm_head.forward(&x)?;
357        logits.to_dtype(DType::F32)
358    }
359
360    pub fn load(vb: VarBuilder, cfg: Config) -> Result<Self> {
361        let wte = embedding(cfg.vocab_size, cfg.dim, vb.pp("model.embed_tokens"))?;
362        let lm_head = linear(cfg.dim, cfg.vocab_size, vb.pp("lm_head"))?;
363        let ln_f = rms_norm(cfg.dim, cfg.norm_eps, vb.pp("model.norm"))?;
364        let blocks: Vec<_> = (0..cfg.n_layers)
365            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), &cfg).unwrap())
366            .collect();
367        Ok(Self {
368            wte,
369            blocks,
370            ln_f,
371            lm_head,
372            config: cfg,
373        })
374    }
375}
candle_transformers/models/llama2_c.rs

candle_transformers/models/
llama2_c.rs
