From 11d4a3c588e9ecb5bcc5a57773973b48036155f0 Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Wed, 24 Apr 2024 09:48:13 +0200 Subject: [PATCH] Add the phi-3 model. (#2120) * Add the phi-3 model. * Faster rope. * Bugfix. * Fix the detokenization. --- candle-examples/examples/phi/main.rs | 81 ++++-- candle-transformers/src/models/mod.rs | 1 + candle-transformers/src/models/phi3.rs | 329 +++++++++++++++++++++++++ 3 files changed, 391 insertions(+), 20 deletions(-) create mode 100644 candle-transformers/src/models/phi3.rs diff --git a/candle-examples/examples/phi/main.rs b/candle-examples/examples/phi/main.rs index 39f4fd58..b65a803d 100644 --- a/candle-examples/examples/phi/main.rs +++ b/candle-examples/examples/phi/main.rs @@ -7,11 +7,13 @@ extern crate accelerate_src; use anyhow::{Error as E, Result}; use clap::{Parser, ValueEnum}; +use candle_examples::token_output_stream::TokenOutputStream; use candle_transformers::models::mixformer::{Config, MixFormerSequentialForCausalLM as MixFormer}; use candle_transformers::models::phi::{Config as PhiConfig, Model as Phi}; +use candle_transformers::models::phi3::{Config as Phi3Config, Model as Phi3}; use candle_transformers::models::quantized_mixformer::MixFormerSequentialForCausalLM as QMixFormer; -use candle::{DType, Device, Tensor}; +use candle::{DType, Device, IndexOp, Tensor}; use candle_nn::VarBuilder; use candle_transformers::generation::LogitsProcessor; use hf_hub::{api::sync::Api, Repo, RepoType}; @@ -20,13 +22,14 @@ use tokenizers::Tokenizer; enum Model { MixFormer(MixFormer), Phi(Phi), + Phi3(Phi3), Quantized(QMixFormer), } struct TextGeneration { model: Model, device: Device, - tokenizer: Tokenizer, + tokenizer: TokenOutputStream, logits_processor: LogitsProcessor, repeat_penalty: f32, repeat_last_n: usize, @@ -49,7 +52,7 @@ impl TextGeneration { let logits_processor = LogitsProcessor::new(seed, temp, top_p); Self { model, - tokenizer, + tokenizer: TokenOutputStream::new(tokenizer), logits_processor, repeat_penalty, repeat_last_n, @@ -61,7 +64,11 @@ impl TextGeneration { fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> { use std::io::Write; println!("starting the inference loop"); - let tokens = self.tokenizer.encode(prompt, true).map_err(E::msg)?; + let tokens = self + .tokenizer + .tokenizer() + .encode(prompt, true) + .map_err(E::msg)?; if tokens.is_empty() { anyhow::bail!("Empty prompts are not supported in the phi model.") } @@ -73,13 +80,14 @@ impl TextGeneration { } let mut tokens = tokens.get_ids().to_vec(); let mut generated_tokens = 0usize; - let eos_token = match self.tokenizer.get_vocab(true).get("<|endoftext|>") { - Some(token) => *token, + let eos_token = match self.tokenizer.get_token("<|endoftext|>") { + Some(token) => token, None => anyhow::bail!("cannot find the endoftext token"), }; print!("{prompt}"); std::io::stdout().flush()?; let start_gen = std::time::Instant::now(); + let mut pos = 0; for index in 0..sample_len { let context_size = if index > 0 { 1 } else { tokens.len() }; let ctxt = &tokens[tokens.len().saturating_sub(context_size)..]; @@ -88,6 +96,7 @@ impl TextGeneration { Model::MixFormer(m) => m.forward(&input)?, Model::Phi(m) => m.forward(&input)?, Model::Quantized(m) => m.forward(&input)?, + Model::Phi3(m) => m.forward(&input, pos)?.i((.., 0, ..))?, }; let logits = logits.squeeze(0)?.to_dtype(DType::F32)?; let logits = if self.repeat_penalty == 1. { @@ -107,9 +116,11 @@ impl TextGeneration { if next_token == eos_token { break; } - let token = self.tokenizer.decode(&[next_token], true).map_err(E::msg)?; - print!("{token}"); - std::io::stdout().flush()?; + if let Some(t) = self.tokenizer.next_token(next_token)? { + print!("{t}"); + std::io::stdout().flush()?; + } + pos += context_size; } let dt = start_gen.elapsed(); println!( @@ -128,6 +139,8 @@ enum WhichModel { V1_5, #[value(name = "2")] V2, + #[value(name = "3")] + V3, #[value(name = "2-old")] V2Old, PuffinPhiV2, @@ -236,6 +249,7 @@ fn main() -> Result<()> { WhichModel::V1 => "microsoft/phi-1".to_string(), WhichModel::V1_5 => "microsoft/phi-1_5".to_string(), WhichModel::V2 | WhichModel::V2Old => "microsoft/phi-2".to_string(), + WhichModel::V3 => "microsoft/Phi-3-mini-4k-instruct".to_string(), WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { "lmz/candle-quantized-phi".to_string() } @@ -253,9 +267,10 @@ fn main() -> Result<()> { WhichModel::V1 => "refs/pr/8".to_string(), WhichModel::V1_5 => "refs/pr/73".to_string(), WhichModel::V2Old => "834565c23f9b28b96ccbeabe614dd906b6db551a".to_string(), - WhichModel::V2 | WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { - "main".to_string() - } + WhichModel::V2 + | WhichModel::V3 + | WhichModel::PuffinPhiV2 + | WhichModel::PhiHermes => "main".to_string(), } } } @@ -264,9 +279,11 @@ 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 | WhichModel::V2Old => { - repo.get("tokenizer.json")? - } + WhichModel::V1 + | WhichModel::V1_5 + | WhichModel::V2 + | WhichModel::V2Old + | WhichModel::V3 => repo.get("tokenizer.json")?, WhichModel::PuffinPhiV2 | WhichModel::PhiHermes => { repo.get("tokenizer-puffin-phi-v2.json")? } @@ -282,14 +299,19 @@ fn main() -> Result<()> { WhichModel::V2 | WhichModel::V2Old => 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")?], + WhichModel::V3 => anyhow::bail!( + "use the quantized or quantized-phi examples for quantized phi-v3" + ), } } else { match args.model { WhichModel::V1 | WhichModel::V1_5 => vec![repo.get("model.safetensors")?], - WhichModel::V2 | WhichModel::V2Old => candle_examples::hub_load_safetensors( - &repo, - "model.safetensors.index.json", - )?, + WhichModel::V2 | WhichModel::V2Old | WhichModel::V3 => { + candle_examples::hub_load_safetensors( + &repo, + "model.safetensors.index.json", + )? + } WhichModel::PuffinPhiV2 => vec![repo.get("model-puffin-phi-v2.safetensors")?], WhichModel::PhiHermes => vec![repo.get("model-phi-hermes-1_3B.safetensors")?], } @@ -306,6 +328,9 @@ fn main() -> Result<()> { WhichModel::V2 | WhichModel::V2Old => Config::v2(), WhichModel::PuffinPhiV2 => Config::puffin_phi_v2(), WhichModel::PhiHermes => Config::phi_hermes_1_3b(), + WhichModel::V3 => { + panic!("use the quantized or quantized-phi examples for quantized phi-v3") + } }; let device = candle_examples::device(args.cpu)?; let model = if args.quantized { @@ -320,7 +345,12 @@ fn main() -> Result<()> { }; Model::Quantized(model) } else { - let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, DType::F32, &device)? }; + let dtype = if args.model == WhichModel::V3 && device.is_cuda() { + DType::BF16 + } else { + DType::F32 + }; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? }; match args.model { WhichModel::V1 | WhichModel::V1_5 | WhichModel::V2 => { let config_filename = repo.get("config.json")?; @@ -329,6 +359,13 @@ fn main() -> Result<()> { let phi = Phi::new(&config, vb)?; Model::Phi(phi) } + WhichModel::V3 => { + let config_filename = repo.get("config.json")?; + let config = std::fs::read_to_string(config_filename)?; + let config: Phi3Config = serde_json::from_str(&config)?; + let phi3 = Phi3::new(&config, vb)?; + Model::Phi3(phi3) + } WhichModel::V2Old => { let config = config(); Model::MixFormer(MixFormer::new_v2(&config, vb)?) @@ -421,6 +458,10 @@ fn mmlu>( m.clear_kv_cache(); m.forward(&input)? } + Model::Phi3(m) => { + m.clear_kv_cache(); + m.forward(&input, 0)? + } Model::Quantized(m) => { m.clear_kv_cache(); m.forward(&input)? diff --git a/candle-transformers/src/models/mod.rs b/candle-transformers/src/models/mod.rs index 5f1a40ad..19c16696 100644 --- a/candle-transformers/src/models/mod.rs +++ b/candle-transformers/src/models/mod.rs @@ -28,6 +28,7 @@ pub mod moondream; pub mod mpt; pub mod persimmon; pub mod phi; +pub mod phi3; pub mod quantized_blip; pub mod quantized_blip_text; pub mod quantized_llama; diff --git a/candle-transformers/src/models/phi3.rs b/candle-transformers/src/models/phi3.rs new file mode 100644 index 00000000..d305e175 --- /dev/null +++ b/candle-transformers/src/models/phi3.rs @@ -0,0 +1,329 @@ +// This implementation is based on: +// https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/modeling_phi3.py +use crate::models::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle::{DType, Device, Module, Result, Tensor, D}; +use candle_nn::VarBuilder; +use std::sync::Arc; + +// https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/config.json +#[derive(Debug, Clone, serde::Deserialize)] +pub struct Config { + pub vocab_size: usize, + pub hidden_act: candle_nn::Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, +} + +impl Config { + 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(dtype: DType, cfg: &Config, dev: &Device) -> Result { + let dim = cfg.head_dim(); + let max_seq_len = cfg.max_position_embeddings; + let inv_freq: Vec<_> = (0..dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) 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)?; + 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 q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } +} + +#[derive(Debug, Clone)] +struct Attention { + qkv_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + rotary_emb: Arc, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl Attention { + fn new(rotary_emb: Arc, cfg: &Config, vb: VarBuilder) -> Result { + let num_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let head_dim = cfg.head_dim(); + let op_size = num_heads * head_dim + 2 * num_kv_heads * head_dim; + let qkv_proj = linear(cfg.hidden_size, op_size, vb.pp("qkv_proj"))?; + let o_proj = linear(num_heads * head_dim, cfg.hidden_size, vb.pp("o_proj"))?; + Ok(Self { + qkv_proj, + o_proj, + rotary_emb, + kv_cache: None, + num_heads, + num_kv_heads, + num_kv_groups: num_heads / num_kv_heads, + head_dim, + }) + } + + fn forward( + &mut self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + + let qkv = self.qkv_proj.forward(xs)?; + let query_pos = self.num_heads * self.head_dim; + let query_states = qkv.narrow(D::Minus1, 0, query_pos)?; + let key_states = qkv.narrow(D::Minus1, query_pos, self.num_kv_heads * self.head_dim)?; + let value_states = qkv.narrow( + D::Minus1, + query_pos + self.num_kv_heads * self.head_dim, + self.num_kv_heads * self.head_dim, + )?; + + 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 = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?; + let value_states = + crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?; + + let attn_output = { + 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, ()))? + .apply(&self.o_proj) + } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +#[derive(Debug, Clone)] +struct Mlp { + gate_up_proj: Linear, + down_proj: Linear, + act_fn: candle_nn::Activation, + i_size: usize, +} + +impl Mlp { + fn new(cfg: &Config, vb: VarBuilder) -> Result { + let hidden_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let gate_up_proj = linear(hidden_size, 2 * i_size, vb.pp("gate_up_proj"))?; + let down_proj = linear(i_size, hidden_size, vb.pp("down_proj"))?; + Ok(Self { + gate_up_proj, + down_proj, + act_fn: cfg.hidden_act, + i_size, + }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let up_states = xs.apply(&self.gate_up_proj)?; + let gate = up_states.narrow(D::Minus1, 0, self.i_size)?; + let up_states = up_states.narrow(D::Minus1, self.i_size, self.i_size)?; + let up_states = (up_states * gate.apply(&self.act_fn))?; + up_states.apply(&self.down_proj) + } +} + +#[derive(Debug, Clone)] +struct DecoderLayer { + self_attn: Attention, + mlp: Mlp, + 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 mlp = Mlp::new(cfg, vb.pp("mlp"))?; + 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, + mlp, + 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.mlp)?; + residual + xs + } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } +} + +#[derive(Debug, Clone)] +pub struct Model { + embed_tokens: candle_nn::Embedding, + layers: Vec, + norm: RmsNorm, + lm_head: Linear, + 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(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?; + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + device: vb.device().clone(), + dtype: vb.dtype(), + }) + } + + fn prepare_decoder_attention_mask( + &self, + b_size: usize, + tgt_len: usize, + seqlen_offset: usize, + ) -> Result { + let mask: Vec<_> = (0..tgt_len) + .flat_map(|i| (0..tgt_len).map(move |j| if i < j { 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) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +}