From a57a91ed1086f629363e6d872abc990545da7900 Mon Sep 17 00:00:00 2001 From: Laurent Date: Thu, 14 Dec 2023 16:27:28 -0600 Subject: [PATCH 01/13] Add the Mixtral model. --- candle-transformers/src/models/mixtral.rs | 45 +++++++++++++++++++++++ candle-transformers/src/models/mod.rs | 1 + 2 files changed, 46 insertions(+) create mode 100644 candle-transformers/src/models/mixtral.rs diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs new file mode 100644 index 0000000000..3c68f20bdb --- /dev/null +++ b/candle-transformers/src/models/mixtral.rs @@ -0,0 +1,45 @@ +#![allow(unused)] +use crate::models::with_tracing::{linear, Embedding as E, Linear}; +/// Mixtral Model +/// https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py +/// https://mistral.ai/news/mixtral-of-experts/ +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; + +/// https://github.com/huggingface/transformers/blob/1a585c1222a56bcaecc070966d558d4a9d862e83/src/transformers/models/mixtral/configuration_mixtral.py#L113 +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Config { + pub(crate) vocab_size: usize, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) rms_norm_eps: f64, + pub(crate) rope_theta: f64, + pub(crate) sliding_window: usize, + pub(crate) num_experts_per_tok: usize, + pub(crate) num_local_experts: usize, +} + +impl Config { + /// https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json + pub fn v0_1_8x7b() -> Self { + Self { + vocab_size: 32000, + hidden_size: 4096, + intermediate_size: 14336, + num_hidden_layers: 32, + num_key_value_heads: 8, + hidden_act: Activation::Silu, + max_position_embeddings: 32768, + rms_norm_eps: 1e-5, + rope_theta: 1e6, + sliding_window: 4096, + num_experts_per_tok: 2, + num_local_experts: 8, + } + } +} diff --git a/candle-transformers/src/models/mod.rs b/candle-transformers/src/models/mod.rs index a9a5667392..94a3bd5be1 100644 --- a/candle-transformers/src/models/mod.rs +++ b/candle-transformers/src/models/mod.rs @@ -14,6 +14,7 @@ pub mod llama2_c_weights; pub mod marian; pub mod mistral; pub mod mixformer; +pub mod mixtral; pub mod mpt; pub mod persimmon; pub mod quantized_blip; From 965bb16b6b7cb95dfd5c35064e0aa744088afeb1 Mon Sep 17 00:00:00 2001 From: Laurent Date: Thu, 14 Dec 2023 16:51:40 -0600 Subject: [PATCH 02/13] Add more of the mixtral layers. --- candle-transformers/src/models/mixtral.rs | 244 +++++++++++++++++++++- 1 file changed, 243 insertions(+), 1 deletion(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 3c68f20bdb..a9d41d6973 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -1,11 +1,12 @@ #![allow(unused)] -use crate::models::with_tracing::{linear, Embedding as E, Linear}; +use crate::models::with_tracing::{linear, linear_no_bias, Embedding as E, Linear}; /// Mixtral Model /// https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py /// https://mistral.ai/news/mixtral-of-experts/ use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; use candle_nn::{Activation, VarBuilder}; use serde::Deserialize; +use std::sync::Arc; /// https://github.com/huggingface/transformers/blob/1a585c1222a56bcaecc070966d558d4a9d862e83/src/transformers/models/mixtral/configuration_mixtral.py#L113 #[derive(Debug, Clone, PartialEq, Deserialize)] @@ -14,6 +15,7 @@ pub struct Config { pub(crate) hidden_size: usize, pub(crate) intermediate_size: usize, pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, pub(crate) num_key_value_heads: usize, pub(crate) hidden_act: Activation, pub(crate) max_position_embeddings: usize, @@ -22,6 +24,7 @@ pub struct Config { pub(crate) sliding_window: usize, pub(crate) num_experts_per_tok: usize, pub(crate) num_local_experts: usize, + pub(crate) use_flash_attn: bool, } impl Config { @@ -32,6 +35,7 @@ impl Config { hidden_size: 4096, intermediate_size: 14336, num_hidden_layers: 32, + num_attention_heads: 32, num_key_value_heads: 8, hidden_act: Activation::Silu, max_position_embeddings: 32768, @@ -40,6 +44,244 @@ impl Config { sliding_window: 4096, num_experts_per_tok: 2, num_local_experts: 8, + use_flash_attn: false, } } } + +#[derive(Debug, Clone)] +struct RmsNorm { + inner: candle_nn::RmsNorm, + span: tracing::Span, +} + +impl RmsNorm { + fn new(size: usize, eps: f64, vb: VarBuilder) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "rms-norm"); + let inner = candle_nn::rms_norm(size, eps, vb)?; + Ok(Self { inner, span }) + } +} + +impl Module for RmsNorm { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + self.inner.forward(x) + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +fn rotate_half(xs: &Tensor) -> Result { + let last_dim = xs.dim(D::Minus1)?; + let xs1 = xs.narrow(D::Minus1, 0, last_dim / 2)?; + let xs2 = xs.narrow(D::Minus1, last_dim / 2, last_dim - last_dim / 2)?; + Tensor::cat(&[&xs2.neg()?, &xs1], D::Minus1) +} + +impl RotaryEmbedding { + fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.hidden_size / cfg.num_attention_heads; + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .collect(); + let inv_freq_len = inv_freq.len(); + let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; + let t = Tensor::arange(0u32, max_seq_len as u32, dev)? + .to_dtype(dtype)? + .reshape((max_seq_len, 1))?; + let freqs = t.matmul(&inv_freq)?; + let freqs = Tensor::cat(&[&freqs, &freqs], D::Minus1)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb_qkv( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offset: usize, + ) -> Result<(Tensor, Tensor)> { + let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; + let cos = self.cos.narrow(0, seqlen_offset, seq_len)?; + let sin = self.sin.narrow(0, seqlen_offset, seq_len)?; + let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let sin = sin.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, seq_len, dim) + let q_embed = (q.broadcast_mul(&cos)? + rotate_half(q)?.broadcast_mul(&sin))?; + let k_embed = (k.broadcast_mul(&cos)? + rotate_half(k)?.broadcast_mul(&sin))?; + Ok((q_embed, k_embed)) + } +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + hidden_size: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, + use_flash_attn: bool, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let num_kv_groups = num_heads / num_kv_heads; + let head_dim = hidden_sz / num_heads; + let q_proj = linear_no_bias(hidden_sz, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear_no_bias(hidden_sz, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear_no_bias(num_heads * head_dim, hidden_sz, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads, + num_kv_heads, + num_kv_groups, + head_dim, + hidden_size: hidden_sz, + rotary_emb, + kv_cache: None, + use_flash_attn: cfg.use_flash_attn, + }) + } + + fn repeat_kv(&self, xs: Tensor) -> Result { + let n_rep = self.num_kv_groups; + if n_rep == 1 { + Ok(xs) + } else { + let (b_sz, num_kv_heads, seq_len, head_dim) = xs.dims4()?; + xs.unsqueeze(2)? + .expand((b_sz, num_kv_heads, n_rep, seq_len, head_dim))? + .reshape((b_sz, num_kv_heads * n_rep, seq_len, head_dim)) + } + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (query_states, key_states) = + self.rotary_emb + .apply_rotary_emb_qkv(&query_states, &key_states, seqlen_offset)?; + + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + let key_states = self.repeat_kv(key_states)?; + let value_states = self.repeat_kv(value_states)?; + + let attn_output = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, q_len > 1)?.transpose(1, 2)? + } else { + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_weights = (query_states.matmul(&key_states.transpose(2, 3)?)? * scale)?; + + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + attn_weights.matmul(&value_states)? + }; + attn_output + .transpose(1, 2)? + .reshape((b_sz, q_len, self.hidden_size))? + .apply(&self.o_proj) + } +} + +#[derive(Debug, Clone)] +struct BlockSparseTop2MLP { + w1: Linear, + w2: Linear, + w3: Linear, + act_fn: Activation, +} + +impl BlockSparseTop2MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_sz = cfg.hidden_size; + let intermediate_sz = cfg.intermediate_size; + let w1 = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("w1"))?; + let w2 = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("w2"))?; + let w3 = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("w3"))?; + Ok(Self { + w1, + w2, + w3, + act_fn: cfg.hidden_act, + }) + } + + fn forward(&self, xs: &Tensor, routing_weights: &Tensor) -> Result { + let lhs = xs.apply(&self.w1)?.apply(&self.act_fn)?; + let rhs = xs.apply(&self.w3)?; + (lhs * rhs)?.apply(&self.w3)? * routing_weights + } +} From 4dcdb2d6dc9ce584267895d75ed65dd45923a0fb Mon Sep 17 00:00:00 2001 From: Laurent Date: Thu, 14 Dec 2023 17:14:15 -0600 Subject: [PATCH 03/13] Add the final layers for mixtral. --- candle-transformers/src/models/mixtral.rs | 153 ++++++++++++++++++++++ 1 file changed, 153 insertions(+) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index a9d41d6973..7b59387fcb 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -285,3 +285,156 @@ impl BlockSparseTop2MLP { (lhs * rhs)?.apply(&self.w3)? * routing_weights } } + +#[derive(Debug, Clone)] +struct SparseMoeBlock { + gate: Linear, + experts: Vec, +} + +impl SparseMoeBlock { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let gate = linear_no_bias(cfg.hidden_size, cfg.num_local_experts, vb.pp("gate"))?; + let mut experts = Vec::with_capacity(cfg.num_local_experts); + let vb = vb.pp("experts"); + for idx in 0..cfg.num_local_experts { + let expert = BlockSparseTop2MLP::new(cfg, vb.pp(idx))?; + experts.push(expert) + } + Ok(SparseMoeBlock { gate, experts }) + } +} + +impl Module for SparseMoeBlock { + fn forward(&self, _xs: &Tensor) -> Result { + todo!() + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + block_sparse_moe: SparseMoeBlock, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl DecoderLayer { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let self_attn = Attention::new(rotary_emb, cfg, vb.pp("self_attn"))?; + let block_sparse_moe = SparseMoeBlock::new(cfg, vb.pp("block_sparse_moe"))?; + let input_layernorm = + RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let post_attention_layernorm = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + self_attn, + block_sparse_moe, + input_layernorm, + post_attention_layernorm, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let residual = xs; + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, attention_mask, seqlen_offset)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = xs + .apply(&self.post_attention_layernorm)? + .apply(&self.block_sparse_moe)?; + residual + xs + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + sliding_window: usize, + device: Device, + dtype: DType, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result { + let vb_m = vb.pp("model"); + let embed_tokens = + candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb_m.device())?); + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_l = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(rotary_emb.clone(), cfg, vb_l.pp(layer_idx))?; + layers.push(layer) + } + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb_m.pp("norm"))?; + let lm_head = linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + sliding_window: cfg.sliding_window, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + // Sliding window mask? + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| { + (0..tgt_len).map(move |j| { + if i < j || j + self.sliding_window < i { + f32::NEG_INFINITY + } else { + 0. + } + }) + }) + .collect(); + let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; + let mask = if seqlen_offset > 0 { + let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, &self.device)?; + Tensor::cat(&[&mask0, &mask], D::Minus1)? + } else { + mask + }; + mask.expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))? + .to_dtype(self.dtype) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let (b_size, seq_len) = input_ids.dims2()?; + let attention_mask = if seq_len <= 1 { + None + } else { + let mask = self.prepare_decoder_attention_mask(b_size, seq_len, seqlen_offset)?; + Some(mask) + }; + let mut xs = self.embed_tokens.forward(input_ids)?; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)? + } + xs.narrow(1, seq_len - 1, 1)? + .apply(&self.norm)? + .apply(&self.lm_head) + } +} From a6159d1d37f9483aea81de7a67edc39001a0349b Mon Sep 17 00:00:00 2001 From: Laurent Date: Thu, 14 Dec 2023 17:34:35 -0600 Subject: [PATCH 04/13] Sketch the expert selection. --- candle-transformers/src/models/mixtral.rs | 31 +++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 7b59387fcb..8e5a163e52 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -306,8 +306,35 @@ impl SparseMoeBlock { } impl Module for SparseMoeBlock { - fn forward(&self, _xs: &Tensor) -> Result { - todo!() + fn forward(&self, xs: &Tensor) -> Result { + let (b_size, seq_len, hidden_dim) = xs.dims3()?; + let xs = xs.reshape(((), hidden_dim))?; + let router_logits = xs.apply(&self.gate)?; + let routing_weights = candle_nn::ops::softmax_last_dim(&router_logits)?; + // TODO: extract topk + // routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + // routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + // expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + let mut ys = xs.zeros_like()?; + for (expert_idx, expert_layer) in self.experts.iter().enumerate() { + // idx, top_x = torch.where(expert_mask[expert_idx]) + // if top_x.shape[0] == 0: + // continue + let top_x = xs.zeros_like()?; // TODO + + // Index the correct hidden states and compute the expert hidden state for + // the current expert. We need to make sure to multiply the output hidden + // states by `routing_weights` on the corresponding tokens (top-1 and top-2) + let current_state = xs.index_select(&top_x, 1)?.reshape(((), hidden_dim))?; + // current_hidden_states = expert_layer(current_state, routing_weights[top_x_list, idx_list, None]) + let current_hidden_states = expert_layer.forward(¤t_state, &routing_weights)?; + + ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; + } + + let xs = xs.reshape((b_size, seq_len, hidden_dim))?; + Ok(xs) } } From 1ff4c62969f8e273a50189c61abeea3ec33a788c Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 09:43:31 -0600 Subject: [PATCH 05/13] Add some expert routing logic. --- candle-transformers/src/models/mixtral.rs | 45 +++++++++++++++++------ 1 file changed, 34 insertions(+), 11 deletions(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 8e5a163e52..54b85cfa77 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -278,11 +278,13 @@ impl BlockSparseTop2MLP { act_fn: cfg.hidden_act, }) } +} - fn forward(&self, xs: &Tensor, routing_weights: &Tensor) -> Result { +impl Module for BlockSparseTop2MLP { + fn forward(&self, xs: &Tensor) -> Result { let lhs = xs.apply(&self.w1)?.apply(&self.act_fn)?; let rhs = xs.apply(&self.w3)?; - (lhs * rhs)?.apply(&self.w3)? * routing_weights + (lhs * rhs)?.apply(&self.w3) } } @@ -290,6 +292,7 @@ impl BlockSparseTop2MLP { struct SparseMoeBlock { gate: Linear, experts: Vec, + num_experts_per_tok: usize, } impl SparseMoeBlock { @@ -301,7 +304,11 @@ impl SparseMoeBlock { let expert = BlockSparseTop2MLP::new(cfg, vb.pp(idx))?; experts.push(expert) } - Ok(SparseMoeBlock { gate, experts }) + Ok(SparseMoeBlock { + gate, + experts, + num_experts_per_tok: cfg.num_experts_per_tok, + }) } } @@ -311,24 +318,40 @@ impl Module for SparseMoeBlock { let xs = xs.reshape(((), hidden_dim))?; let router_logits = xs.apply(&self.gate)?; let routing_weights = candle_nn::ops::softmax_last_dim(&router_logits)?; - // TODO: extract topk + + // In order to extract topk, we extract the data from the tensor and manipulate it + // directly. Maybe we will want to use some custom ops instead at some point. + let routing_weights = routing_weights.to_dtype(DType::F32)?.to_vec2::()?; + // routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + let mut top_k = Vec::with_capacity(routing_weights.len()); + let mut top_x = vec![vec![]; self.experts.len()]; + for (row_idx, rw) in routing_weights.iter().enumerate() { + let mut dst = (0..rw.len() as u32).collect::>(); + dst.sort_by(|&i, &j| rw[j as usize].total_cmp(&rw[i as usize])); + dst.truncate(self.num_experts_per_tok); + for expert_idx in dst.iter() { + top_x[*expert_idx as usize].push(row_idx as u32) + } + top_k.push(dst); + } + // routing_weights /= routing_weights.sum(dim=-1, keepdim=True) // expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) let mut ys = xs.zeros_like()?; for (expert_idx, expert_layer) in self.experts.iter().enumerate() { - // idx, top_x = torch.where(expert_mask[expert_idx]) - // if top_x.shape[0] == 0: - // continue - let top_x = xs.zeros_like()?; // TODO - + let top_x = &top_x[expert_idx]; + if top_x.is_empty() { + continue; + } + let top_x = Tensor::new(top_x.as_slice(), xs.device())?; // Index the correct hidden states and compute the expert hidden state for // the current expert. We need to make sure to multiply the output hidden // states by `routing_weights` on the corresponding tokens (top-1 and top-2) - let current_state = xs.index_select(&top_x, 1)?.reshape(((), hidden_dim))?; + let current_state = xs.index_select(&top_x, 0)?.reshape(((), hidden_dim))?; // current_hidden_states = expert_layer(current_state, routing_weights[top_x_list, idx_list, None]) - let current_hidden_states = expert_layer.forward(¤t_state, &routing_weights)?; + let current_hidden_states = expert_layer.forward(¤t_state)?; ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; } From 944f2c77c7613976988a8a7c02279a19edc508fe Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 09:50:10 -0600 Subject: [PATCH 06/13] Hopefully finish the routing logic for mixtral. --- candle-transformers/src/models/mixtral.rs | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 54b85cfa77..2d36f09314 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -324,16 +324,24 @@ impl Module for SparseMoeBlock { let routing_weights = routing_weights.to_dtype(DType::F32)?.to_vec2::()?; // routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) - let mut top_k = Vec::with_capacity(routing_weights.len()); + // top_x contains the row indexes to evaluate for each expert. let mut top_x = vec![vec![]; self.experts.len()]; + let mut selected_rws = vec![vec![]; self.experts.len()]; for (row_idx, rw) in routing_weights.iter().enumerate() { let mut dst = (0..rw.len() as u32).collect::>(); dst.sort_by(|&i, &j| rw[j as usize].total_cmp(&rw[i as usize])); - dst.truncate(self.num_experts_per_tok); - for expert_idx in dst.iter() { - top_x[*expert_idx as usize].push(row_idx as u32) + let mut sum_routing_weights = 0f32; + for &expert_idx in dst.iter().take(self.num_experts_per_tok) { + let expert_idx = expert_idx as usize; + let routing_weight = rw[expert_idx]; + sum_routing_weights += routing_weight; + top_x[expert_idx].push(row_idx as u32); + } + for &expert_idx in dst.iter().take(self.num_experts_per_tok) { + let expert_idx = expert_idx as usize; + let routing_weight = rw[expert_idx]; + selected_rws[expert_idx].push(routing_weight / sum_routing_weights) } - top_k.push(dst); } // routing_weights /= routing_weights.sum(dim=-1, keepdim=True) @@ -346,13 +354,14 @@ impl Module for SparseMoeBlock { continue; } let top_x = Tensor::new(top_x.as_slice(), xs.device())?; + let selected_rws = Tensor::new(selected_rws[expert_idx].as_slice(), xs.device())?; // Index the correct hidden states and compute the expert hidden state for // the current expert. We need to make sure to multiply the output hidden // states by `routing_weights` on the corresponding tokens (top-1 and top-2) let current_state = xs.index_select(&top_x, 0)?.reshape(((), hidden_dim))?; // current_hidden_states = expert_layer(current_state, routing_weights[top_x_list, idx_list, None]) let current_hidden_states = expert_layer.forward(¤t_state)?; - + let current_hidden_states = current_hidden_states.broadcast_mul(&selected_rws)?; ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; } From 5442f779be16928b31f35e1758f0767a5e04ef59 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 11:13:40 -0600 Subject: [PATCH 07/13] Add the mixtral example. --- candle-examples/examples/mixtral/main.rs | 263 ++++++++++++++++++++++ candle-transformers/src/models/mixtral.rs | 4 +- 2 files changed, 265 insertions(+), 2 deletions(-) create mode 100644 candle-examples/examples/mixtral/main.rs diff --git a/candle-examples/examples/mixtral/main.rs b/candle-examples/examples/mixtral/main.rs new file mode 100644 index 0000000000..b1d71bde72 --- /dev/null +++ b/candle-examples/examples/mixtral/main.rs @@ -0,0 +1,263 @@ +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use anyhow::{Error as E, Result}; +use clap::Parser; + +use candle_transformers::models::mixtral::{Config, Model}; + +use candle::{DType, Device, Tensor}; +use candle_examples::token_output_stream::TokenOutputStream; +use candle_nn::VarBuilder; +use candle_transformers::generation::LogitsProcessor; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use tokenizers::Tokenizer; + +struct TextGeneration { + model: Model, + device: Device, + tokenizer: TokenOutputStream, + logits_processor: LogitsProcessor, + repeat_penalty: f32, + repeat_last_n: usize, +} + +impl TextGeneration { + #[allow(clippy::too_many_arguments)] + fn new( + model: Model, + tokenizer: Tokenizer, + seed: u64, + temp: Option, + top_p: Option, + repeat_penalty: f32, + repeat_last_n: usize, + device: &Device, + ) -> Self { + let logits_processor = LogitsProcessor::new(seed, temp, top_p); + Self { + model, + tokenizer: TokenOutputStream::new(tokenizer), + logits_processor, + repeat_penalty, + repeat_last_n, + device: device.clone(), + } + } + + fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> { + use std::io::Write; + self.tokenizer.clear(); + let mut tokens = self + .tokenizer + .tokenizer() + .encode(prompt, true) + .map_err(E::msg)? + .get_ids() + .to_vec(); + for &t in tokens.iter() { + if let Some(t) = self.tokenizer.next_token(t)? { + print!("{t}") + } + } + std::io::stdout().flush()?; + + let mut generated_tokens = 0usize; + let eos_token = match self.tokenizer.get_token("") { + Some(token) => token, + None => anyhow::bail!("cannot find the token"), + }; + let start_gen = std::time::Instant::now(); + for index in 0..sample_len { + let context_size = if index > 0 { 1 } else { tokens.len() }; + let start_pos = tokens.len().saturating_sub(context_size); + let ctxt = &tokens[start_pos..]; + let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?; + let logits = self.model.forward(&input, start_pos)?; + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let logits = if self.repeat_penalty == 1. { + logits + } else { + let start_at = tokens.len().saturating_sub(self.repeat_last_n); + candle_transformers::utils::apply_repeat_penalty( + &logits, + self.repeat_penalty, + &tokens[start_at..], + )? + }; + + let next_token = self.logits_processor.sample(&logits)?; + tokens.push(next_token); + generated_tokens += 1; + if next_token == eos_token { + break; + } + if let Some(t) = self.tokenizer.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + } + let dt = start_gen.elapsed(); + if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? { + print!("{rest}"); + } + std::io::stdout().flush()?; + println!( + "\n{generated_tokens} tokens generated ({:.2} token/s)", + generated_tokens as f64 / dt.as_secs_f64(), + ); + Ok(()) + } +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// Run on CPU rather than on GPU. + #[arg(long)] + cpu: bool, + + /// Enable tracing (generates a trace-timestamp.json file). + #[arg(long)] + tracing: bool, + + #[arg(long)] + use_flash_attn: bool, + + #[arg(long)] + prompt: String, + + /// The temperature used to generate samples. + #[arg(long)] + temperature: Option, + + /// Nucleus sampling probability cutoff. + #[arg(long)] + top_p: Option, + + /// The seed to use when generating random samples. + #[arg(long, default_value_t = 299792458)] + seed: u64, + + /// The length of the sample to generate (in tokens). + #[arg(long, short = 'n', default_value_t = 100)] + sample_len: usize, + + #[arg(long, default_value = "mistralai/Mixtral-8x7B-v0.1")] + model_id: String, + + #[arg(long, default_value = "main")] + revision: String, + + #[arg(long)] + tokenizer_file: Option, + + #[arg(long)] + weight_files: Option, + + /// Penalty to be applied for repeating tokens, 1. means no penalty. + #[arg(long, default_value_t = 1.1)] + repeat_penalty: f32, + + /// The context size to consider for the repeat penalty. + #[arg(long, default_value_t = 64)] + repeat_last_n: usize, +} + +fn main() -> Result<()> { + use tracing_chrome::ChromeLayerBuilder; + use tracing_subscriber::prelude::*; + + let args = Args::parse(); + let _guard = if args.tracing { + let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); + tracing_subscriber::registry().with(chrome_layer).init(); + Some(guard) + } else { + None + }; + println!( + "avx: {}, neon: {}, simd128: {}, f16c: {}", + candle::utils::with_avx(), + candle::utils::with_neon(), + candle::utils::with_simd128(), + candle::utils::with_f16c() + ); + println!( + "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}", + args.temperature.unwrap_or(0.), + args.repeat_penalty, + args.repeat_last_n + ); + + let start = std::time::Instant::now(); + let api = Api::new()?; + let repo = api.repo(Repo::with_revision( + args.model_id, + RepoType::Model, + args.revision, + )); + let tokenizer_filename = match args.tokenizer_file { + Some(file) => std::path::PathBuf::from(file), + None => repo.get("tokenizer.json")?, + }; + let filenames = match args.weight_files { + Some(files) => files + .split(',') + .map(std::path::PathBuf::from) + .collect::>(), + None => { + vec![ + repo.get("pytorch_model-00001-of-00019.safetensors")?, + repo.get("pytorch_model-00002-of-00019.safetensors")?, + repo.get("pytorch_model-00003-of-00019.safetensors")?, + repo.get("pytorch_model-00004-of-00019.safetensors")?, + repo.get("pytorch_model-00005-of-00019.safetensors")?, + repo.get("pytorch_model-00006-of-00019.safetensors")?, + repo.get("pytorch_model-00007-of-00019.safetensors")?, + repo.get("pytorch_model-00008-of-00019.safetensors")?, + repo.get("pytorch_model-00009-of-00019.safetensors")?, + repo.get("pytorch_model-00010-of-00019.safetensors")?, + repo.get("pytorch_model-00011-of-00019.safetensors")?, + repo.get("pytorch_model-00012-of-00019.safetensors")?, + repo.get("pytorch_model-00013-of-00019.safetensors")?, + repo.get("pytorch_model-00014-of-00019.safetensors")?, + repo.get("pytorch_model-00015-of-00019.safetensors")?, + repo.get("pytorch_model-00016-of-00019.safetensors")?, + repo.get("pytorch_model-00017-of-00019.safetensors")?, + repo.get("pytorch_model-00018-of-00019.safetensors")?, + repo.get("pytorch_model-00019-of-00019.safetensors")?, + ] + } + }; + println!("retrieved the files in {:?}", start.elapsed()); + let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?; + + let start = std::time::Instant::now(); + let config = Config::v0_1_8x7b(args.use_flash_attn); + let device = candle_examples::device(args.cpu)?; + let dtype = if device.is_cuda() { + DType::BF16 + } else { + DType::F32 + }; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? }; + let model = Model::new(&config, vb)?; + println!("loaded the model in {:?}", start.elapsed()); + + let mut pipeline = TextGeneration::new( + model, + tokenizer, + args.seed, + args.temperature, + args.top_p, + args.repeat_penalty, + args.repeat_last_n, + &device, + ); + pipeline.run(&args.prompt, args.sample_len)?; + Ok(()) +} diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 2d36f09314..f608b916d3 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -29,7 +29,7 @@ pub struct Config { impl Config { /// https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json - pub fn v0_1_8x7b() -> Self { + pub fn v0_1_8x7b(use_flash_attn: bool) -> Self { Self { vocab_size: 32000, hidden_size: 4096, @@ -44,7 +44,7 @@ impl Config { sliding_window: 4096, num_experts_per_tok: 2, num_local_experts: 8, - use_flash_attn: false, + use_flash_attn, } } } From 74f3dbb61888f13526fb5f05ecb535e17fc7ba21 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 11:38:23 -0600 Subject: [PATCH 08/13] Fix the weight filenames. --- candle-examples/examples/mixtral/main.rs | 38 ++++++++++++------------ 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/candle-examples/examples/mixtral/main.rs b/candle-examples/examples/mixtral/main.rs index b1d71bde72..fcde03c14f 100644 --- a/candle-examples/examples/mixtral/main.rs +++ b/candle-examples/examples/mixtral/main.rs @@ -211,25 +211,25 @@ fn main() -> Result<()> { .collect::>(), None => { vec![ - repo.get("pytorch_model-00001-of-00019.safetensors")?, - repo.get("pytorch_model-00002-of-00019.safetensors")?, - repo.get("pytorch_model-00003-of-00019.safetensors")?, - repo.get("pytorch_model-00004-of-00019.safetensors")?, - repo.get("pytorch_model-00005-of-00019.safetensors")?, - repo.get("pytorch_model-00006-of-00019.safetensors")?, - repo.get("pytorch_model-00007-of-00019.safetensors")?, - repo.get("pytorch_model-00008-of-00019.safetensors")?, - repo.get("pytorch_model-00009-of-00019.safetensors")?, - repo.get("pytorch_model-00010-of-00019.safetensors")?, - repo.get("pytorch_model-00011-of-00019.safetensors")?, - repo.get("pytorch_model-00012-of-00019.safetensors")?, - repo.get("pytorch_model-00013-of-00019.safetensors")?, - repo.get("pytorch_model-00014-of-00019.safetensors")?, - repo.get("pytorch_model-00015-of-00019.safetensors")?, - repo.get("pytorch_model-00016-of-00019.safetensors")?, - repo.get("pytorch_model-00017-of-00019.safetensors")?, - repo.get("pytorch_model-00018-of-00019.safetensors")?, - repo.get("pytorch_model-00019-of-00019.safetensors")?, + repo.get("model-00001-of-00019.safetensors")?, + repo.get("model-00002-of-00019.safetensors")?, + repo.get("model-00003-of-00019.safetensors")?, + repo.get("model-00004-of-00019.safetensors")?, + repo.get("model-00005-of-00019.safetensors")?, + repo.get("model-00006-of-00019.safetensors")?, + repo.get("model-00007-of-00019.safetensors")?, + repo.get("model-00008-of-00019.safetensors")?, + repo.get("model-00009-of-00019.safetensors")?, + repo.get("model-00010-of-00019.safetensors")?, + repo.get("model-00011-of-00019.safetensors")?, + repo.get("model-00012-of-00019.safetensors")?, + repo.get("model-00013-of-00019.safetensors")?, + repo.get("model-00014-of-00019.safetensors")?, + repo.get("model-00015-of-00019.safetensors")?, + repo.get("model-00016-of-00019.safetensors")?, + repo.get("model-00017-of-00019.safetensors")?, + repo.get("model-00018-of-00019.safetensors")?, + repo.get("model-00019-of-00019.safetensors")?, ] } }; From 9175dcc95e374b4828f4c2b55cddaa75e29586d4 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 13:00:06 -0600 Subject: [PATCH 09/13] Bugfix. --- candle-transformers/src/models/mixtral.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index f608b916d3..72be2c1848 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -284,7 +284,7 @@ impl Module for BlockSparseTop2MLP { fn forward(&self, xs: &Tensor) -> Result { let lhs = xs.apply(&self.w1)?.apply(&self.act_fn)?; let rhs = xs.apply(&self.w3)?; - (lhs * rhs)?.apply(&self.w3) + (lhs * rhs)?.apply(&self.w2) } } From 8ba6cdbf3b70d0fd12e5e912dd08903b34f69663 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 13:41:55 -0600 Subject: [PATCH 10/13] Another fix. --- candle-transformers/src/models/mixtral.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index 72be2c1848..e01dfb590a 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -89,7 +89,7 @@ impl RotaryEmbedding { let max_seq_len = cfg.max_position_embeddings; let inv_freq: Vec<_> = (0..dim) .step_by(2) - .map(|i| 1f32 / 10000f32.powf(i as f32 / dim as f32)) + .map(|i| 1f32 / (cfg.rope_theta as f32).powf(i as f32 / dim as f32)) .collect(); let inv_freq_len = inv_freq.len(); let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(dtype)?; From 74fc82c5114cc1a1faedcc06c3120951ae2751c4 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 13:51:55 -0600 Subject: [PATCH 11/13] Yet another fix + remove the unused pragma. --- candle-transformers/src/models/mixtral.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index e01dfb590a..dd39ca60eb 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -1,9 +1,8 @@ -#![allow(unused)] -use crate::models::with_tracing::{linear, linear_no_bias, Embedding as E, Linear}; +use crate::models::with_tracing::{linear_no_bias, Linear}; /// Mixtral Model /// https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py /// https://mistral.ai/news/mixtral-of-experts/ -use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle::{DType, Device, Module, Result, Tensor, D}; use candle_nn::{Activation, VarBuilder}; use serde::Deserialize; use std::sync::Arc; @@ -365,8 +364,8 @@ impl Module for SparseMoeBlock { ys = ys.index_add(&top_x, ¤t_hidden_states, 0)?; } - let xs = xs.reshape((b_size, seq_len, hidden_dim))?; - Ok(xs) + let ys = ys.reshape((b_size, seq_len, hidden_dim))?; + Ok(ys) } } From 5471437ffff345811754643d9f360ed8fc48f6c7 Mon Sep 17 00:00:00 2001 From: Laurent Date: Fri, 15 Dec 2023 13:54:03 -0600 Subject: [PATCH 12/13] Shape fix. --- candle-transformers/src/models/mixtral.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/candle-transformers/src/models/mixtral.rs b/candle-transformers/src/models/mixtral.rs index dd39ca60eb..ede74d3f9b 100644 --- a/candle-transformers/src/models/mixtral.rs +++ b/candle-transformers/src/models/mixtral.rs @@ -353,7 +353,8 @@ impl Module for SparseMoeBlock { continue; } let top_x = Tensor::new(top_x.as_slice(), xs.device())?; - let selected_rws = Tensor::new(selected_rws[expert_idx].as_slice(), xs.device())?; + let selected_rws = + Tensor::new(selected_rws[expert_idx].as_slice(), xs.device())?.reshape(((), 1))?; // Index the correct hidden states and compute the expert hidden state for // the current expert. We need to make sure to multiply the output hidden // states by `routing_weights` on the corresponding tokens (top-1 and top-2) From 3c9db2a76095560d4df6f31a56d31029e0b56306 Mon Sep 17 00:00:00 2001 From: laurent Date: Fri, 15 Dec 2023 20:17:25 +0000 Subject: [PATCH 13/13] Add a readme. --- candle-examples/examples/mixtral/README.md | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 candle-examples/examples/mixtral/README.md diff --git a/candle-examples/examples/mixtral/README.md b/candle-examples/examples/mixtral/README.md new file mode 100644 index 0000000000..aec5c14807 --- /dev/null +++ b/candle-examples/examples/mixtral/README.md @@ -0,0 +1,25 @@ +# candle-mixtral: 8x7b LLM using a sparse mixture of experts. + +Mixtral-8x7B-v0.1 is a pretrained generative LLM with 56 billion parameters. + +- [Blog post](https://mistral.ai/news/mixtral-of-experts/) from Mistral announcing the model release. +- [Model card](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) on the HuggingFace Hub. + +## Running the example + +```bash +$ cargo run --example mixtral --release -- --prompt "def print_prime(n): " +def print_prime(n): # n is the number of prime numbers to be printed + i = 2 + count = 0 + while (count < n): + if (isPrime(i)): + print(i) + count += 1 + i += 1 + +def isPrime(n): + for x in range(2, int(n**0.5)+1): + if (n % x == 0): + ... +```