diff options
Diffstat (limited to 'candle-examples/examples/llama/model.rs')
| -rw-r--r-- | candle-examples/examples/llama/model.rs | 184 |
1 files changed, 104 insertions, 80 deletions
diff --git a/candle-examples/examples/llama/model.rs b/candle-examples/examples/llama/model.rs index ce2e6d2e..f3e30ec9 100644 --- a/candle-examples/examples/llama/model.rs +++ b/candle-examples/examples/llama/model.rs @@ -12,6 +12,7 @@ pub struct Config { pub n_layer: usize, pub n_head: usize, pub n_embd: usize, + pub n_key_value_head: usize, } impl Config { @@ -23,6 +24,7 @@ impl Config { n_layer: 32, n_head: 32, n_embd: 4096, + n_key_value_head: 32, } } } @@ -33,17 +35,37 @@ pub struct Cache { pub use_kv_cache: bool, #[allow(clippy::type_complexity)] kvs: Arc<Mutex<Vec<Option<(Tensor, Tensor)>>>>, + cos: Tensor, + sin: Tensor, device: Device, } impl Cache { - pub fn new(use_kv_cache: bool, config: &Config, device: &Device) -> Self { - Self { + pub fn new(use_kv_cache: bool, config: &Config, device: &Device) -> Result<Self> { + // precompute freqs_cis + let n_elem = config.n_embd / config.n_head; + let theta: Vec<_> = (0..n_elem) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / n_elem as f32)) + .collect(); + let theta = Tensor::new(theta.as_slice(), device)?; + let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, device)? + .to_dtype(DType::F32)? + .reshape((MAX_SEQ_LEN, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + // This is different from the paper, see: + // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 + let idx_theta = Tensor::cat(&[&idx_theta, &idx_theta], D::Minus1)?; + let cos = idx_theta.cos()?; + let sin = idx_theta.sin()?; + Ok(Self { masks: Arc::new(Mutex::new(HashMap::new())), use_kv_cache, kvs: Arc::new(Mutex::new(vec![None; config.n_layer])), device: device.clone(), - } + cos, + sin, + }) } fn mask(&self, t: usize) -> Result<Tensor> { @@ -97,7 +119,7 @@ impl RmsNorm { let (b_sz, seq_len, hidden_size) = x.shape().r3()?; let norm_x = (x.sqr()?.sum_keepdim(2)? / hidden_size as f64)?; let norm_x = norm_x.broadcast_as((b_sz, seq_len, hidden_size))?; - let x_normed = (x / (norm_x + 1e-5)?.sqrt()?)?; + let x_normed = (x / (norm_x + 1e-6)?.sqrt()?)?; let size = self.scale.shape().r1()?; let scale = self .scale @@ -110,63 +132,52 @@ impl RmsNorm { } struct CausalSelfAttention { - c_attn: Linear, - c_proj: Linear, + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, n_head: usize, + n_key_value_head: usize, + head_dim: usize, cache: Cache, } impl CausalSelfAttention { - fn new(c_attn: Linear, c_proj: Linear, n_head: usize, cache: &Cache) -> Self { - Self { - c_attn, - c_proj, - n_head, - cache: cache.clone(), - } - } - - fn apply_rotary_emb(&self, x: &Tensor, freqs_cis: &Tensor) -> Result<Tensor> { - let mut dims = x.dims().to_vec(); - let fcis_dims = freqs_cis.dims(); - let freqs_cis = if dims[2] < fcis_dims[1] { - freqs_cis.narrow(1, 0, dims[2])? - } else { - freqs_cis.clone() - }; - let v = dims.pop().unwrap(); - dims.push(v / 2); - dims.push(2); - let x = x.reshape(dims)?; - let re_x = x.narrow(D::Minus1, 0, 1)?; - let im_x = x.narrow(D::Minus1, 1, 1)?; - let re_f = freqs_cis - .narrow(D::Minus1, 0, 1)? - .broadcast_as(re_x.shape())?; - let im_f = freqs_cis - .narrow(D::Minus1, 1, 1)? - .broadcast_as(im_x.shape())?; - let re = ((&re_x * &re_f)? - (&im_x * &im_f)?)?; - let im = ((&re_x * &im_f)? + (&im_x * &re_f)?)?; - let rope = Tensor::cat(&[&re, &im], D::Minus1)?; - let rope = rope.flatten_from(D::Minus2)?; + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize) -> Result<Tensor> { + let (b_sz, _, seq_len, n_embd) = x.shape().r4()?; + let cos = self.cache.cos.narrow(0, index_pos, seq_len)?; + let sin = self.cache.sin.narrow(0, index_pos, seq_len)?; + let cos = cos.broadcast_as((b_sz, 1, seq_len, n_embd))?; + let sin = sin.broadcast_as((b_sz, 1, seq_len, n_embd))?; + let x1 = x.narrow(D::Minus1, 0, n_embd / 2)?; + let x2 = x.narrow(D::Minus1, n_embd / 2, n_embd / 2)?; + let rotate_x = Tensor::cat(&[&x2.neg()?, &x1], D::Minus1)?; + let rope = (x.broadcast_mul(&cos)? + rotate_x.broadcast_mul(&sin)?)?; Ok(rope) } - fn forward(&self, x: &Tensor, freqs_cis: &Tensor, block_idx: usize) -> Result<Tensor> { + fn forward(&self, x: &Tensor, index_pos: usize, block_idx: usize) -> Result<Tensor> { let x_dtype = x.dtype(); let (b_sz, seq_len, n_embd) = x.shape().r3()?; - let qkv = self.c_attn.forward(x)?; - let qkv = qkv.to_dtype(DType::F32)?; - let q = qkv.narrow(D::Minus1, 0, n_embd)?; - let k = qkv.narrow(D::Minus1, n_embd, n_embd)?; - let v = qkv.narrow(D::Minus1, 2 * n_embd, n_embd)?; - let target_dim = [b_sz, seq_len, self.n_head, n_embd / self.n_head]; - let k = k.reshape(target_dim.as_slice())?.transpose(1, 2)?; - let q = q.reshape(target_dim.as_slice())?.transpose(1, 2)?; - let mut v = v.reshape(target_dim.as_slice())?.transpose(1, 2)?; - let q = self.apply_rotary_emb(&q, freqs_cis)?; - let mut k = self.apply_rotary_emb(&k, freqs_cis)?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.n_head, self.head_dim))? + .transpose(1, 2)? + .to_dtype(DType::F32)?; + let k = k + .reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))? + .transpose(1, 2)? + .to_dtype(DType::F32)?; + let mut v = v + .reshape((b_sz, seq_len, self.n_key_value_head, self.head_dim))? + .transpose(1, 2)? + .to_dtype(DType::F32)?; + + let q = self.apply_rotary_emb(&q, index_pos)?; + let mut k = self.apply_rotary_emb(&k, index_pos)?; if self.cache.use_kv_cache { let mut cache = self.cache.kvs.lock().unwrap(); @@ -189,7 +200,9 @@ impl CausalSelfAttention { cache[block_idx] = Some((k.clone(), v.clone())) } - let att = (q.matmul(&k.t()?)? / (k.dim(D::Minus1)? as f64).sqrt())?; + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; let mask = self.cache.mask(seq_len)?.broadcast_as(att.shape())?; let att = masked_fill(&att, &mask, f32::NEG_INFINITY)?; let att = att.softmax(D::Minus1)?; @@ -197,31 +210,42 @@ impl CausalSelfAttention { let y = att.matmul(&v.contiguous()?)?; let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, n_embd])?; let y = y.to_dtype(x_dtype)?; - let y = self.c_proj.forward(&y)?; + let y = self.o_proj.forward(&y)?; Ok(y) } + fn repeat_kv(&self, x: Tensor) -> Result<Tensor> { + let n_rep = self.n_head / self.n_key_value_head; + if n_rep == 1 { + Ok(x) + } else { + let (b_sz, n_kv_head, seq_len, head_dim) = x.shape().r4()?; + let x = x + .unsqueeze(2)? + .expand((b_sz, n_kv_head, n_rep, seq_len, head_dim))? + .reshape((b_sz, n_kv_head, n_rep, seq_len, head_dim))?; + Ok(x) + } + } + fn load(vb: VarBuilder, cache: &Cache, cfg: &Config) -> Result<Self> { let size_in = cfg.hidden_size; - let size = (cfg.hidden_size / cfg.n_head) * cfg.n_head; - let q_proj = vb.get((size_in, size), "q_proj.weight")?; - let k_proj = vb.get((size_in, size), "k_proj.weight")?; - let v_proj = vb.get((size_in, size), "v_proj.weight")?; - // Invert the transformation from: - // https://github.com/huggingface/transformers/blob/2642d8d04b14c18199ebe7b35f976da02df61752/src/transformers/models/llama/convert_llama_weights_to_hf.py#L101 - let n_head = cfg.n_head; - let q_proj = q_proj - .reshape((n_head, 2, size / n_head / 2, size_in))? - .transpose(1, 2)? - .reshape((size_in, size))?; - let k_proj = k_proj - .reshape((n_head, 2, size / n_head / 2, size_in))? - .transpose(1, 2)? - .reshape((size_in, size))?; - let attn_weight = Tensor::cat(&[q_proj, k_proj, v_proj], 0)?; - let c_attn = Linear::new(attn_weight, None); - let o_proj = linear(size, size_in, vb.pp("o_proj"))?; - Ok(Self::new(c_attn, o_proj, cfg.n_head, cache)) + let size_q = (cfg.hidden_size / cfg.n_head) * cfg.n_head; + let size_kv = (cfg.hidden_size / cfg.n_head) * cfg.n_key_value_head; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + n_head: cfg.n_head, + n_key_value_head: cfg.n_key_value_head, + head_dim: cfg.hidden_size / cfg.n_head, + cache: cache.clone(), + }) } } @@ -279,12 +303,12 @@ impl Block { } } - fn forward(&self, x: &Tensor, freqs_cis: &Tensor, block_idx: usize) -> Result<Tensor> { - let x = (self - .attn - .forward(&self.rms_1.forward(x)?, freqs_cis, block_idx)? - + x)?; - let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + x)?; + fn forward(&self, x: &Tensor, index_pos: usize, block_idx: usize) -> Result<Tensor> { + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; Ok(x) } @@ -320,11 +344,11 @@ impl Llama { } } - pub fn forward(&self, x: &Tensor, freqs_cis: &Tensor) -> Result<Tensor> { + pub fn forward(&self, x: &Tensor, index_pos: usize) -> Result<Tensor> { let (_b_sz, seq_len) = x.shape().r2()?; let mut x = self.wte.forward(x)?; for (block_idx, block) in self.blocks.iter().enumerate() { - x = block.forward(&x, freqs_cis, block_idx)?; + x = block.forward(&x, index_pos, block_idx)?; } let x = self.ln_f.forward(&x)?; let x = x.i((.., seq_len - 1, ..))?; |