huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support for MQA for llama v2. (#205)

Browse Source 
* Support for MQA for llama v2.

* More llama-v2.

* Move the rotary embedding precomputation in the cache.

* Add a v2 flag.

* Use the hf model.
This commit is contained in:
Laurent Mazare
2023-07-20 07:39:04 +02:00
committed by GitHub
parent c34f932319
commit 12d6dc018d
3 changed files with 123 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

184
candle-examples/examples/llama/model.rs
View File

@ -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, ..))?;