diff options
| -rw-r--r-- | candle-examples/examples/phi/main.rs | 46 | ||||
| -rw-r--r-- | candle-nn/src/activation.rs | 1 | ||||
| -rw-r--r-- | candle-transformers/src/models/mod.rs | 1 | ||||
| -rw-r--r-- | candle-transformers/src/models/phi.rs | 365 |
4 files changed, 402 insertions, 11 deletions
diff --git a/candle-examples/examples/phi/main.rs b/candle-examples/examples/phi/main.rs index c5c7de28..ea99c706 100644 --- a/candle-examples/examples/phi/main.rs +++ b/candle-examples/examples/phi/main.rs @@ -8,6 +8,7 @@ use anyhow::{Error as E, Result}; use clap::{Parser, ValueEnum}; use candle_transformers::models::mixformer::{Config, MixFormerSequentialForCausalLM as MixFormer}; +use candle_transformers::models::phi::{Config as PhiConfig, Model as Phi}; use candle_transformers::models::quantized_mixformer::MixFormerSequentialForCausalLM as QMixFormer; use candle::{DType, Device, Tensor}; @@ -18,6 +19,7 @@ use tokenizers::Tokenizer; enum Model { MixFormer(MixFormer), + Phi(Phi), Quantized(QMixFormer), } @@ -84,6 +86,7 @@ impl TextGeneration { let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?; let logits = match &mut self.model { Model::MixFormer(m) => m.forward(&input)?, + Model::Phi(m) => m.forward(&input)?, Model::Quantized(m) => m.forward(&input)?, }; let logits = logits.squeeze(0)?.to_dtype(DType::F32)?; @@ -117,7 +120,7 @@ impl TextGeneration { } } -#[derive(Clone, Copy, Debug, ValueEnum)] +#[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq)] enum WhichModel { #[value(name = "1")] V1, @@ -125,6 +128,9 @@ enum WhichModel { V1_5, #[value(name = "2")] V2, + // TODO: Make this the default once it has been battle tested. + #[value(name = "2-new")] + V2New, PuffinPhiV2, PhiHermes, } @@ -230,7 +236,7 @@ fn main() -> Result<()> { match args.model { WhichModel::V1 => "microsoft/phi-1".to_string(), WhichModel::V1_5 => "microsoft/phi-1_5".to_string(), - WhichModel::V2 => "microsoft/phi-2".to_string(), + WhichModel::V2 | WhichModel::V2New => "microsoft/phi-2".to_string(), WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { "lmz/candle-quantized-phi".to_string() } @@ -248,7 +254,9 @@ fn main() -> Result<()> { WhichModel::V1 => "refs/pr/2".to_string(), WhichModel::V1_5 => "refs/pr/18".to_string(), WhichModel::V2 => "834565c23f9b28b96ccbeabe614dd906b6db551a".to_string(), - WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => "main".to_string(), + WhichModel::V2New | WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { + "main".to_string() + } } } } @@ -257,7 +265,9 @@ fn main() -> Result<()> { let tokenizer_filename = match args.tokenizer { Some(file) => std::path::PathBuf::from(file), None => match args.model { - WhichModel::V1 | WhichModel::V1_5 | WhichModel::V2 => repo.get("tokenizer.json")?, + WhichModel::V1 | WhichModel::V1_5 | WhichModel::V2 | WhichModel::V2New => { + repo.get("tokenizer.json")? + } WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { repo.get("tokenizer-puffin-phi-v2.json")? } @@ -270,14 +280,14 @@ fn main() -> Result<()> { match args.model { WhichModel::V1 => vec![repo.get("model-v1-q4k.gguf")?], WhichModel::V1_5 => vec![repo.get("model-q4k.gguf")?], - WhichModel::V2 => vec![repo.get("model-v2-q4k.gguf")?], + WhichModel::V2 | WhichModel::V2New => vec![repo.get("model-v2-q4k.gguf")?], WhichModel::PuffinPhiV2 => vec![repo.get("model-puffin-phi-v2-q4k.gguf")?], WhichModel::PhiHermes => vec![repo.get("model-phi-hermes-1_3B-q4k.gguf")?], } } else { match args.model { WhichModel::V1 | WhichModel::V1_5 => vec![repo.get("model.safetensors")?], - WhichModel::V2 => candle_examples::hub_load_safetensors( + WhichModel::V2 | WhichModel::V2New => candle_examples::hub_load_safetensors( &repo, "model.safetensors.index.json", )?, @@ -291,25 +301,35 @@ fn main() -> Result<()> { let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?; let start = std::time::Instant::now(); - let config = match args.model { + let config = || match args.model { WhichModel::V1 => Config::v1(), WhichModel::V1_5 => Config::v1_5(), - WhichModel::V2 => Config::v2(), + WhichModel::V2 | WhichModel::V2New => Config::v2(), WhichModel::PuffinPhiV2 => Config::puffin_phi_v2(), WhichModel::PhiHermes => Config::phi_hermes_1_3b(), }; - let (model, device) = if args.quantized { + let (model, device) = if args.model == WhichModel::V2New { + let device = candle_examples::device(args.cpu)?; + let config_filename = repo.get("config.json")?; + let config = std::fs::read_to_string(config_filename)?; + let config: PhiConfig = serde_json::from_str(&config)?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, DType::F32, &device)? }; + let phi = Phi::new(&config, vb)?; + (Model::Phi(phi), device) + } else if args.quantized { let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(&filenames[0])?; + let config = config(); let model = match args.model { - WhichModel::V2 => QMixFormer::new_v2(&config, vb)?, + WhichModel::V2 | WhichModel::V2New => QMixFormer::new_v2(&config, vb)?, _ => QMixFormer::new(&config, vb)?, }; (Model::Quantized(model), Device::Cpu) } else { let device = candle_examples::device(args.cpu)?; + let config = config(); let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, DType::F32, &device)? }; let model = match args.model { - WhichModel::V2 => MixFormer::new_v2(&config, vb)?, + WhichModel::V2 | WhichModel::V2New => MixFormer::new_v2(&config, vb)?, _ => MixFormer::new(&config, vb)?, }; (Model::MixFormer(model), device) @@ -392,6 +412,10 @@ fn mmlu<P: AsRef<std::path::Path>>( m.clear_kv_cache(); m.forward(&input)? } + Model::Phi(m) => { + m.clear_kv_cache(); + m.forward(&input)? + } Model::Quantized(m) => { m.clear_kv_cache(); m.forward(&input)? diff --git a/candle-nn/src/activation.rs b/candle-nn/src/activation.rs index 80b750ed..e00463f0 100644 --- a/candle-nn/src/activation.rs +++ b/candle-nn/src/activation.rs @@ -6,6 +6,7 @@ use serde::Deserialize; pub enum Activation { #[default] Gelu, + #[serde(alias = "gelu_new")] NewGelu, Relu, Relu2, diff --git a/candle-transformers/src/models/mod.rs b/candle-transformers/src/models/mod.rs index a60b5a06..9af6df69 100644 --- a/candle-transformers/src/models/mod.rs +++ b/candle-transformers/src/models/mod.rs @@ -17,6 +17,7 @@ pub mod mixformer; pub mod mixtral; pub mod mpt; pub mod persimmon; +pub mod phi; pub mod quantized_blip; pub mod quantized_blip_text; pub mod quantized_llama; diff --git a/candle-transformers/src/models/phi.rs b/candle-transformers/src/models/phi.rs new file mode 100644 index 00000000..a635f3ce --- /dev/null +++ b/candle-transformers/src/models/phi.rs @@ -0,0 +1,365 @@ +use crate::models::with_tracing::{layer_norm, linear, Embedding, LayerNorm, Linear}; +/// Phi model. +/// https://huggingface.co/microsoft/phi-2 +/// There is an alternative implementation of the phi model in mixformers.rs. +/// This corresponds to the model update made with the following commit: +/// https://huggingface.co/microsoft/phi-2/commit/cb2f4533604d8b67de604e7df03bfe6f3ca22869 +use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; +use candle_nn::{Activation, VarBuilder}; +use serde::Deserialize; + +// https://huggingface.co/microsoft/phi-2/blob/main/configuration_phi.py +#[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_attention_heads: usize, + pub(crate) num_key_value_heads: Option<usize>, + pub(crate) hidden_act: Activation, + pub(crate) max_position_embeddings: usize, + pub(crate) layer_norm_eps: f64, + pub(crate) tie_word_embeddings: bool, + pub(crate) rope_theta: f32, + pub(crate) partial_rotary_factor: f64, + pub(crate) qk_layernorm: bool, +} + +impl Config { + fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } + + fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } +} + +#[derive(Debug, Clone)] +struct RotaryEmbedding { + sin: Tensor, + cos: Tensor, +} + +impl RotaryEmbedding { + fn new(cfg: &Config, dev: &Device) -> Result<Self> { + let dim = (cfg.partial_rotary_factor * cfg.head_dim() as f64) as usize; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.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)?; + let t = Tensor::arange(0u32, cfg.max_position_embeddings as u32, dev)? + .to_dtype(DType::F32)? + .reshape((cfg.max_position_embeddings, 1))?; + let freqs = t.matmul(&inv_freq)?; + Ok(Self { + sin: freqs.sin()?, + cos: freqs.cos()?, + }) + } + + fn apply_rotary_emb(&self, xs: &Tensor, seqlen_offset: usize) -> Result<Tensor> { + let (_b_size, seqlen, _, _headdim) = xs.dims4()?; + let (_rotary_seqlen, rotary_dim) = self.cos.dims2()?; + let rotary_dim = rotary_dim * 2; + let xs_rot = xs.i((.., .., .., ..rotary_dim))?; + let xs_pass = xs.i((.., .., .., rotary_dim..))?; + let xs12 = xs_rot.chunk(2, D::Minus1)?; + let (xs1, xs2) = (&xs12[0], &xs12[1]); + let c = self.cos.narrow(0, seqlen_offset, seqlen)?.unsqueeze(1)?; + let s = self.sin.narrow(0, seqlen_offset, seqlen)?.unsqueeze(1)?; + let xs_rot = Tensor::cat( + &[ + (xs1.broadcast_mul(&c)? - xs2.broadcast_mul(&s)?)?, + (xs1.broadcast_mul(&s)? + xs2.broadcast_mul(&c)?)?, + ], + D::Minus1, + )?; + Tensor::cat(&[&xs_rot, &xs_pass], D::Minus1) + } +} + +#[derive(Debug, Clone)] +#[allow(clippy::upper_case_acronyms)] +struct MLP { + fc1: Linear, + fc2: Linear, + act: Activation, +} + +impl MLP { + fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> { + let fc1 = linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?; + Ok(Self { + fc1, + fc2, + act: cfg.hidden_act, + }) + } +} + +impl Module for MLP { + fn forward(&self, xs: &Tensor) -> Result<Tensor> { + xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2) + } +} + +#[derive(Clone)] +struct Attention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + dense: Linear, + kv_cache: Option<(Tensor, Tensor)>, + q_layernorm: Option<LayerNorm>, + k_layernorm: Option<LayerNorm>, + rotary_emb: RotaryEmbedding, + softmax_scale: f64, + num_heads: usize, + num_kv_heads: usize, + head_dim: usize, + span: tracing::Span, +} + +fn get_mask(size: usize, device: &Device) -> Result<Tensor> { + let mask: Vec<_> = (0..size) + .flat_map(|i| (0..size).map(move |j| u8::from(j > i))) + .collect(); + Tensor::from_slice(&mask, (size, size), device) +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> { + let shape = mask.shape(); + let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?; + let m = mask.where_cond(&on_true, on_false)?; + Ok(m) +} + +impl Attention { + fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> { + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads(); + let head_dim = cfg.head_dim(); + let q_proj = linear(cfg.hidden_size, num_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(cfg.hidden_size, num_kv_heads * head_dim, vb.pp("v_proj"))?; + let dense = linear(num_heads * head_dim, cfg.hidden_size, vb.pp("dense"))?; + // Alternative rope scalings are not supported. + let rotary_emb = RotaryEmbedding::new(cfg, vb.device())?; + let (q_layernorm, k_layernorm) = if cfg.qk_layernorm { + let q_layernorm = layer_norm(head_dim, cfg.layer_norm_eps, vb.pp("q_layernorm"))?; + let k_layernorm = layer_norm(head_dim, cfg.layer_norm_eps, vb.pp("k_layernorm"))?; + (Some(q_layernorm), Some(k_layernorm)) + } else { + (None, None) + }; + let softmax_scale = 1f64 / (head_dim as f64).sqrt(); + Ok(Self { + q_proj, + k_proj, + v_proj, + dense, + kv_cache: None, + q_layernorm, + k_layernorm, + rotary_emb, + softmax_scale, + num_heads, + num_kv_heads, + head_dim, + span: tracing::span!(tracing::Level::TRACE, "attention"), + }) + } + + fn repeat_kv(&self, xs: Tensor) -> Result<Tensor> { + let n_rep = self.num_heads / self.num_kv_heads; + 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, mask: Option<&Tensor>) -> Result<Tensor> { + let _enter = self.span.enter(); + let (b_size, seq_len, _n_embd) = 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 = match &self.q_layernorm { + None => query_states, + Some(ln) => query_states.apply(ln)?, + }; + let key_states = match &self.k_layernorm { + None => key_states, + Some(ln) => key_states.apply(ln)?, + }; + + let query_states = query_states + .reshape((b_size, seq_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_size, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_size, seq_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + // Rotary embeddings. + let seqlen_offset = match &self.kv_cache { + None => 0, + Some((prev_k, _)) => prev_k.dim(1)?, + }; + let query_states = self + .rotary_emb + .apply_rotary_emb(&query_states, seqlen_offset)?; + let key_states = self + .rotary_emb + .apply_rotary_emb(&key_states, seqlen_offset)?; + + // KV cache. + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let k = Tensor::cat(&[prev_k, &key_states], 2)?; + let v = Tensor::cat(&[prev_v, &value_states], 2)?; + (k, v) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + // Repeat kv. + let key_states = self.repeat_kv(key_states)?.contiguous()?; + let value_states = self.repeat_kv(value_states)?.contiguous()?; + + let attn_weights = (query_states + .to_dtype(DType::F32)? + .contiguous()? + .matmul(&key_states.to_dtype(DType::F32)?.t()?)? + * self.softmax_scale)?; + let attn_weights = match mask { + None => attn_weights, + Some(mask) => masked_fill( + &attn_weights, + &mask.broadcast_left((b_size, self.num_heads))?, + f32::NEG_INFINITY, + )?, + }; + let attn_weights = + candle_nn::ops::softmax_last_dim(&attn_weights)?.to_dtype(value_states.dtype())?; + let attn_output = attn_weights.matmul(&value_states)?; + let attn_output = attn_output + .transpose(1, 2)? + .reshape((b_size, seq_len, ()))?; + attn_output.apply(&self.dense) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: MLP, + input_layernorm: LayerNorm, + span: tracing::Span, +} + +impl DecoderLayer { + fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> { + let self_attn = Attention::new(cfg, vb.pp("self_attn"))?; + let mlp = MLP::new(cfg, vb.pp("mlp"))?; + let input_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb.pp("input_layernorm"), + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + span: tracing::span!(tracing::Level::TRACE, "block"), + }) + } + + fn forward(&mut self, xs: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> { + let _enter = self.span.enter(); + let residual = xs; + let xs = xs.apply(&self.input_layernorm)?; + let attn_outputs = self.self_attn.forward(&xs, mask)?; + let feed_forward_hidden_states = self.mlp.forward(&xs)?; + attn_outputs + feed_forward_hidden_states + residual + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Clone)] +pub struct Model { + embed_tokens: Embedding, + layers: Vec<DecoderLayer>, + final_layernorm: LayerNorm, + lm_head: Linear, + span: tracing::Span, +} + +impl Model { + pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> { + let vb_m = vb.pp("model"); + let embed_tokens = + Embedding::new(cfg.vocab_size, cfg.hidden_size, vb_m.pp("embed_tokens"))?; + let final_layernorm = layer_norm( + cfg.hidden_size, + cfg.layer_norm_eps, + vb_m.pp("final_layernorm"), + )?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + let vb_m = vb_m.pp("layers"); + for layer_idx in 0..cfg.num_hidden_layers { + let layer = DecoderLayer::new(cfg, vb_m.pp(layer_idx))?; + layers.push(layer) + } + let lm_head = linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + final_layernorm, + lm_head, + span: tracing::span!(tracing::Level::TRACE, "model"), + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result<Tensor> { + let _enter = self.span.enter(); + let (_b_size, seq_len) = xs.dims2()?; + let mut xs = xs.apply(&self.embed_tokens)?; + let mask = if seq_len <= 1 { + None + } else { + Some(get_mask(seq_len, xs.device())?) + }; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, mask.as_ref())? + } + xs.apply(&self.final_layernorm)? + .narrow(1, seq_len - 1, 1)? + .apply(&self.lm_head)? + .squeeze(1) + } + + pub fn clear_kv_cache(&mut self) { + self.layers.iter_mut().for_each(|b| b.clear_kv_cache()) + } +} |