From b1d42231fba72ff508bc02c34c65feab5409ec4c Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 07:23:02 +0100 Subject: [PATCH 01/17] Start sketching the whisper model. --- candle-examples/examples/whisper/main.rs | 281 +++++++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100644 candle-examples/examples/whisper/main.rs diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs new file mode 100644 index 00000000..7b2d0c74 --- /dev/null +++ b/candle-examples/examples/whisper/main.rs @@ -0,0 +1,281 @@ +#![allow(dead_code)] +// https://github.com/openai/whisper/blob/main/whisper/model.py + +use anyhow::{Error as E, Result}; +use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor}; +use clap::Parser; +use std::collections::HashMap; + +const DTYPE: DType = DType::F32; + +struct VarBuilder<'a> { + safetensors: Option<(HashMap, Vec>)>, + dtype: DType, + device: Device, +} + +impl<'a> VarBuilder<'a> { + pub fn from_safetensors( + safetensors: Vec>, + dtype: DType, + device: Device, + ) -> Self { + let mut routing = HashMap::new(); + for (index, sf) in safetensors.iter().enumerate() { + for k in sf.names() { + routing.insert(k.to_string(), index); + } + } + Self { + safetensors: Some((routing, safetensors)), + device, + dtype, + } + } + + pub fn zeros(dtype: DType, device: Device) -> Self { + Self { + safetensors: None, + device, + dtype, + } + } + + pub fn get>(&self, s: S, tensor_name: &str) -> candle::Result { + let s: Shape = s.into(); + match &self.safetensors { + None => Tensor::zeros(s, self.dtype, &self.device), + Some((routing, safetensors)) => { + // Unwrap or 0 just to let the proper error flow. + let index = routing.get(tensor_name).unwrap_or(&0); + let tensor = safetensors[*index] + .tensor(tensor_name, &self.device)? + .to_dtype(self.dtype)?; + if *tensor.shape() != s { + let msg = format!("shape mismatch for {tensor_name}"); + Err(candle::Error::UnexpectedShape { + msg, + expected: s, + got: tensor.shape().clone(), + })? + } + Ok(tensor) + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum HiddenAct { + Gelu, + Relu, +} + +impl HiddenAct { + fn forward(&self, xs: &Tensor) -> candle::Result { + match self { + Self::Gelu => xs.gelu(), + Self::Relu => xs.relu(), + } + } +} + +#[derive(Debug, Clone, PartialEq)] +struct Config { + n_mels: usize, + n_audio_ctx: usize, + n_audio_state: usize, + n_audio_head: usize, + n_audio_layer: usize, + n_vocab: usize, + n_text_ctx: usize, + n_text_state: usize, + n_text_head: usize, + n_text_layer: usize, +} + +struct Embedding { + embeddings: Tensor, + hidden_size: usize, +} + +impl Embedding { + fn new(embeddings: Tensor, hidden_size: usize) -> Self { + Self { + embeddings, + hidden_size, + } + } + + fn load(vocab_size: usize, hidden_size: usize, p: &str, vb: &VarBuilder) -> Result { + let embeddings = vb.get((vocab_size, hidden_size), &format!("{p}.weight"))?; + Ok(Self::new(embeddings, hidden_size)) + } + + fn forward(&self, indexes: &Tensor) -> Result { + let mut final_dims = indexes.dims().to_vec(); + final_dims.push(self.hidden_size); + let indexes = indexes.flatten_all()?; + let values = Tensor::embedding(&indexes, &self.embeddings)?; + let values = values.reshape(final_dims)?; + Ok(values) + } +} + +struct Linear { + weight: Tensor, + bias: Option, +} + +impl Linear { + fn load(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result { + let weight = vb.get((size2, size1), &format!("{p}.weight"))?; + let bias = vb.get(size2, &format!("{p}.bias"))?; + Ok(Self { + weight, + bias: Some(bias), + }) + } + + fn load_no_bias(size1: usize, size2: usize, p: &str, vb: &VarBuilder) -> Result { + let weight = vb.get((size2, size1), &format!("{p}.weight"))?; + Ok(Self { weight, bias: None }) + } + + fn forward(&self, x: &Tensor) -> candle::Result { + let (bsize, _, _) = x.shape().r3()?; + let w = self.weight.broadcast_left(bsize)?.t()?; + let x = x.matmul(&w)?; + match &self.bias { + None => Ok(x), + Some(bias) => x.broadcast_add(bias), + } + } +} + +struct Dropout { + pr: f64, +} + +impl Dropout { + fn new(pr: f64) -> Self { + Self { pr } + } + + fn forward(&self, x: &Tensor) -> Result { + // TODO + Ok(x.clone()) + } +} + +// This layer norm version handles both weight and bias so removes the mean. +struct LayerNorm { + weight: Tensor, + bias: Tensor, + eps: f64, +} + +impl LayerNorm { + fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self { + Self { weight, bias, eps } + } + + fn load(size: usize, eps: f64, p: &str, vb: &VarBuilder) -> Result { + let weight = vb.get(size, &format!("{p}.weight"))?; + let bias = vb.get(size, &format!("{p}.bias"))?; + Ok(Self::new(weight, bias, eps)) + } + + fn forward(&self, x: &Tensor) -> Result { + let (_bsize, _seq_len, hidden_size) = x.shape().r3()?; + let mean_x = (x.sum(&[2])? / hidden_size as f64)?; + let x = x.broadcast_sub(&mean_x)?; + let norm_x = ((&x * &x)?.sum(&[2])? / hidden_size as f64)?; + let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?; + let x = x_normed + .broadcast_mul(&self.weight)? + .broadcast_add(&self.bias)?; + Ok(x) + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L62 +struct MultiHeadAttention { + query: Linear, + key: Linear, + value: Linear, + out: Linear, + n_head: usize, +} + +impl MultiHeadAttention { + fn load(n_head: usize, n: usize, p: &str, vb: &VarBuilder) -> Result { + let query = Linear::load(n, n, &format!("{p}.query"), vb)?; + let value = Linear::load_no_bias(n, n, &format!("{p}.value"), vb)?; + let key = Linear::load(n, n, &format!("{p}.key"), vb)?; + let out = Linear::load(n, n, &format!("{p}.out"), vb)?; + Ok(Self { + query, + key, + value, + out, + n_head, + }) + } + + fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> { + let q = self.query.forward(x)?; + let k = self.key.forward(x)?; + let v = self.value.forward(x)?; + let (wv, qk) = self.qkv_attention(&q, &k, &v)?; + let out = self.out.forward(&wv)?; + Ok((out, qk)) + } + + fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)> { + let (n_batch, n_ctx, n_state) = q.shape().r3()?; + let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head]; + let scale = ((n_state / self.n_head) as f64).powf(-0.25); + let q = (q.reshape(target_dims)?.transpose(1, 2)? * scale)?; + let k = (k.reshape(target_dims)?.transpose(1, 2)?.transpose(2, 3)? * scale)?; + let v = v.reshape(target_dims)?.transpose(1, 2)?; + let qk = q.matmul(&k)?; + let w = qk.softmax(qk.rank() - 1)?; + let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?; + let qk = qk.detach()?; + Ok((wv, qk)) + } +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// Run on CPU rather than on GPU. + #[arg(long)] + cpu: bool, + + #[arg(long)] + tokenizer_config: String, + + #[arg(long)] + weights: String, +} + +fn main() -> Result<()> { + use tokenizers::Tokenizer; + + let args = Args::parse(); + let device = if args.cpu { + Device::Cpu + } else { + Device::new_cuda(0)? + }; + + let mut tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?; + let _tokenizer = tokenizer.with_padding(None).with_truncation(None); + + let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? }; + let weights = weights.deserialize()?; + let _vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device); + Ok(()) +} From 0ca2af6940e2812e65a2e2af7f53ce9b8c1980d4 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 07:43:36 +0100 Subject: [PATCH 02/17] Add the residual attention block. --- candle-examples/examples/whisper/main.rs | 92 +++++++++++++++++++----- 1 file changed, 75 insertions(+), 17 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 7b2d0c74..1e32c972 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -176,14 +176,14 @@ struct LayerNorm { } impl LayerNorm { - fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self { - Self { weight, bias, eps } - } - - fn load(size: usize, eps: f64, p: &str, vb: &VarBuilder) -> Result { + fn load(size: usize, p: &str, vb: &VarBuilder) -> Result { let weight = vb.get(size, &format!("{p}.weight"))?; let bias = vb.get(size, &format!("{p}.bias"))?; - Ok(Self::new(weight, bias, eps)) + Ok(Self { + weight, + bias, + eps: 1e-5, + }) } fn forward(&self, x: &Tensor) -> Result { @@ -209,11 +209,11 @@ struct MultiHeadAttention { } impl MultiHeadAttention { - fn load(n_head: usize, n: usize, p: &str, vb: &VarBuilder) -> Result { - let query = Linear::load(n, n, &format!("{p}.query"), vb)?; - let value = Linear::load_no_bias(n, n, &format!("{p}.value"), vb)?; - let key = Linear::load(n, n, &format!("{p}.key"), vb)?; - let out = Linear::load(n, n, &format!("{p}.out"), vb)?; + fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result { + let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?; + let value = Linear::load_no_bias(n_state, n_state, &format!("{p}.value"), vb)?; + let key = Linear::load(n_state, n_state, &format!("{p}.key"), vb)?; + let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?; Ok(Self { query, key, @@ -223,16 +223,16 @@ impl MultiHeadAttention { }) } - fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> { + fn forward(&self, x: &Tensor) -> Result { let q = self.query.forward(x)?; let k = self.key.forward(x)?; let v = self.value.forward(x)?; - let (wv, qk) = self.qkv_attention(&q, &k, &v)?; + let wv = self.qkv_attention(&q, &k, &v)?; let out = self.out.forward(&wv)?; - Ok((out, qk)) + Ok(out) } - fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)> { + fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { let (n_batch, n_ctx, n_state) = q.shape().r3()?; let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head]; let scale = ((n_state / self.n_head) as f64).powf(-0.25); @@ -242,8 +242,66 @@ impl MultiHeadAttention { let qk = q.matmul(&k)?; let w = qk.softmax(qk.rank() - 1)?; let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?; - let qk = qk.detach()?; - Ok((wv, qk)) + Ok(wv) + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L111 +struct ResidualAttentionBlock { + attn: MultiHeadAttention, + attn_ln: LayerNorm, + cross_attn: Option, + cross_attn_ln: Option, + mlp_linear1: Linear, + mlp_linear2: Linear, + mlp_ln: LayerNorm, +} + +impl ResidualAttentionBlock { + fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result { + let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?; + let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?; + let (cross_attn, cross_attn_ln) = if ca { + let cross_attn = + MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?; + let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?; + (Some(cross_attn), Some(cross_attn_ln)) + } else { + (None, None) + }; + let n_mlp = n_state * 4; + let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?; + let mlp_linear2 = Linear::load(n_mlp, n_state, &format!("{p}.mlp.2"), vb)?; + let mlp_ln = LayerNorm::load(n_state, &format!("{p}.mlp_ln"), vb)?; + Ok(Self { + attn, + attn_ln, + cross_attn, + cross_attn_ln, + mlp_linear1, + mlp_linear2, + mlp_ln, + }) + } + + fn forward(&self, x: &Tensor) -> Result { + let attn = self.attn.forward(&self.attn_ln.forward(x)?)?; + let mut x = (x + attn)?; + // Cross-Attn + if let Some(cross_attn) = &self.cross_attn { + x = cross_attn.forward(&x)? + } + if let Some(cross_attn_ln) = &self.cross_attn_ln { + x = cross_attn_ln.forward(&x)? + } + // Mlp + let mlp = self.mlp_linear2.forward( + &self + .mlp_linear1 + .forward(&self.mlp_ln.forward(&x)?)? + .gelu()?, + )?; + Ok((x + mlp)?) } } From 512dd9e4d6c278d72f3942da937a40d6cd45f933 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 08:01:06 +0100 Subject: [PATCH 03/17] Flesh out the whisper example. --- candle-examples/examples/whisper/main.rs | 84 ++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 1e32c972..e0292e17 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -305,6 +305,90 @@ impl ResidualAttentionBlock { } } +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 +struct AudioEncoder { + conv1: Linear, // TODO + conv2: Linear, // TODO + blocks: Vec, + ln_post: LayerNorm, +} + +impl AudioEncoder { + fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { + let n_state = cfg.n_audio_state; + let n_head = cfg.n_audio_head; + let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?; + let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?; + let blocks = (0..cfg.n_audio_layer) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) + }) + .collect::>>()?; + let ln_post = LayerNorm::load(n_state, &format!("{p}.ln_post"), vb)?; + Ok(Self { + conv1, + conv2, + blocks, + ln_post, + }) + } + fn forward(&self, _x: &Tensor) -> Result { + todo!() + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176 +struct TextDecoder { + token_embedding: Embedding, + blocks: Vec, + ln: LayerNorm, + mask: Tensor, +} + +impl TextDecoder { + fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { + let n_state = cfg.n_text_state; + let n_head = cfg.n_text_head; + let token_embedding = + Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?; + let blocks = (0..cfg.n_text_layer) + .map(|i| { + ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) + }) + .collect::>>()?; + let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?; + let mask = Tensor::new(&[0u32], &vb.device)?; + Ok(Self { + token_embedding, + blocks, + ln, + mask, + }) + } + fn forward(&self, _tokens: &Tensor, _enc: &Tensor) -> Result { + todo!() + } +} + +// https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L221 +struct Whisper { + encoder: AudioEncoder, + decoder: TextDecoder, +} + +impl Whisper { + fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { + let encoder = AudioEncoder::load(&format!("{p}.encoder"), vb, cfg)?; + let decoder = TextDecoder::load(&format!("{p}.decoder"), vb, cfg)?; + Ok(Self { encoder, decoder }) + } + fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result { + let enc = self.encoder.forward(mel)?; + let dec = self.decoder.forward(tokens, &enc)?; + Ok(dec) + } +} + #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { From 6728a856765ca23f6bf50cd9246bb2f8e63ee41a Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 08:32:03 +0100 Subject: [PATCH 04/17] Add more to the whisper inference. --- candle-examples/examples/whisper/main.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index e0292e17..3e780a2c 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -1,5 +1,7 @@ #![allow(dead_code)] // https://github.com/openai/whisper/blob/main/whisper/model.py +// TODO: +// - kv-cache support? use anyhow::{Error as E, Result}; use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor}; @@ -309,6 +311,7 @@ impl ResidualAttentionBlock { struct AudioEncoder { conv1: Linear, // TODO conv2: Linear, // TODO + positional_embedding: Tensor, blocks: Vec, ln_post: LayerNorm, } @@ -319,6 +322,7 @@ impl AudioEncoder { let n_head = cfg.n_audio_head; let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?; let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?; + let positional_embedding = Tensor::new(&[0u32], &vb.device)?; // TODO let blocks = (0..cfg.n_audio_layer) .map(|i| { ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) @@ -328,12 +332,21 @@ impl AudioEncoder { Ok(Self { conv1, conv2, + positional_embedding, blocks, ln_post, }) } - fn forward(&self, _x: &Tensor) -> Result { - todo!() + fn forward(&self, x: &Tensor) -> Result { + let x = self.conv1.forward(x)?.gelu()?; + let x = self.conv2.forward(&x)?.gelu()?; + let x = x.transpose(1, 2)?; + let mut x = x.broadcast_add(&self.positional_embedding)?; + for block in self.blocks.iter() { + x = block.forward(&x)? + } + let x = self.ln_post.forward(&x)?; + Ok(x) } } @@ -357,7 +370,7 @@ impl TextDecoder { }) .collect::>>()?; let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?; - let mask = Tensor::new(&[0u32], &vb.device)?; + let mask = Tensor::new(&[0u32], &vb.device)?; // TODO Ok(Self { token_embedding, blocks, From 998cfda9c4baefd0fb3407724c344158356d2d31 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 09:03:25 +0100 Subject: [PATCH 05/17] Sinusoid embeddings. --- candle-examples/examples/whisper/main.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 3e780a2c..1930a7a9 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -307,6 +307,21 @@ impl ResidualAttentionBlock { } } +fn sinusoids(length: usize, channels: usize) -> Result { + let max_timescale = 10000f32; + let log_timescale_increment = max_timescale.ln() / (channels / 2 - 1) as f32; + let inv_timescales: Vec<_> = (0..channels / 2) + .map(|i| (i as f32 * (-log_timescale_increment)).exp()) + .collect(); + let arange: Vec<_> = (0..length).map(|c| c as f32).collect(); + let inv_timescales = Tensor::new(inv_timescales.as_slice(), &Device::Cpu)?.unsqueeze(0)?; + let arange = Tensor::new(arange.as_slice(), &Device::Cpu)?.unsqueeze(1)?; + let sh = (length, channels / 2); + let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?; + let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?; + Ok(sincos) +} + // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 struct AudioEncoder { conv1: Linear, // TODO @@ -322,7 +337,7 @@ impl AudioEncoder { let n_head = cfg.n_audio_head; let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?; let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?; - let positional_embedding = Tensor::new(&[0u32], &vb.device)?; // TODO + let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?; let blocks = (0..cfg.n_audio_layer) .map(|i| { ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) From c09aa4b0f4c1fb074aecbf71aea3e96cadebea60 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 09:14:28 +0100 Subject: [PATCH 06/17] Add more to the forward pass and fix the cross-attention. --- candle-examples/examples/whisper/main.rs | 37 +++++++++++++++++------- 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 1930a7a9..129c73b8 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -225,10 +225,10 @@ impl MultiHeadAttention { }) } - fn forward(&self, x: &Tensor) -> Result { + fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result { let q = self.query.forward(x)?; - let k = self.key.forward(x)?; - let v = self.value.forward(x)?; + let k = self.key.forward(xa.unwrap_or(x))?; + let v = self.value.forward(xa.unwrap_or(x))?; let wv = self.qkv_attention(&q, &k, &v)?; let out = self.out.forward(&wv)?; Ok(out) @@ -286,16 +286,16 @@ impl ResidualAttentionBlock { }) } - fn forward(&self, x: &Tensor) -> Result { - let attn = self.attn.forward(&self.attn_ln.forward(x)?)?; + fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result { + let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?; let mut x = (x + attn)?; // Cross-Attn - if let Some(cross_attn) = &self.cross_attn { - x = cross_attn.forward(&x)? - } if let Some(cross_attn_ln) = &self.cross_attn_ln { x = cross_attn_ln.forward(&x)? } + if let Some(cross_attn) = &self.cross_attn { + x = cross_attn.forward(&x, xa)? + } // Mlp let mlp = self.mlp_linear2.forward( &self @@ -358,7 +358,7 @@ impl AudioEncoder { let x = x.transpose(1, 2)?; let mut x = x.broadcast_add(&self.positional_embedding)?; for block in self.blocks.iter() { - x = block.forward(&x)? + x = block.forward(&x, None)? } let x = self.ln_post.forward(&x)?; Ok(x) @@ -368,6 +368,7 @@ impl AudioEncoder { // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L176 struct TextDecoder { token_embedding: Embedding, + positional_embedding: Tensor, blocks: Vec, ln: LayerNorm, mask: Tensor, @@ -377,8 +378,11 @@ impl TextDecoder { fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { let n_state = cfg.n_text_state; let n_head = cfg.n_text_head; + let n_ctx = cfg.n_text_ctx; let token_embedding = Embedding::load(cfg.n_vocab, n_state, &format!("{p}.token_embedding"), vb)?; + let positional_embedding = + vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?; let blocks = (0..cfg.n_text_layer) .map(|i| { ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) @@ -388,13 +392,24 @@ impl TextDecoder { let mask = Tensor::new(&[0u32], &vb.device)?; // TODO Ok(Self { token_embedding, + positional_embedding, blocks, ln, mask, }) } - fn forward(&self, _tokens: &Tensor, _enc: &Tensor) -> Result { - todo!() + fn forward(&self, x: &Tensor, xa: &Tensor) -> Result { + let x_dims = x.dims(); + let last = x_dims[x_dims.len() - 1]; + let token_embedding = self.token_embedding.forward(x)?; + let positional_embedding = self.positional_embedding.narrow(0, 0, last)?; + let mut x = (token_embedding + positional_embedding)?; + for block in self.blocks.iter() { + x = block.forward(&x, Some(xa))? + } + let x = self.ln.forward(&x)?; + let logits = x.matmul(&self.token_embedding.embeddings.t()?)?; + Ok(logits) } } From d71b31144d286df98d2ae12bde66976794f33640 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 09:29:19 +0100 Subject: [PATCH 07/17] Add a weight extraction script. --- .../examples/whisper/extract_weights.py | 13 +++++++++++++ candle-examples/examples/whisper/main.rs | 17 +++++++++++++++++ 2 files changed, 30 insertions(+) create mode 100644 candle-examples/examples/whisper/extract_weights.py diff --git a/candle-examples/examples/whisper/extract_weights.py b/candle-examples/examples/whisper/extract_weights.py new file mode 100644 index 00000000..d6ccffc6 --- /dev/null +++ b/candle-examples/examples/whisper/extract_weights.py @@ -0,0 +1,13 @@ +# Get the checkpoint from +# https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt + +import torch +from safetensors.torch import save_file + +data = torch.load("tiny.en.pt") +weights = {} +for k, v in data["model_state_dict"].items(): + weights[k] = v.contiguous() + print(k, v.shape) +save_file(weights, "tiny.en.safetensors") +print(data["dims"]) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 129c73b8..dac6df84 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -96,6 +96,23 @@ struct Config { n_text_layer: usize, } +impl Config { + fn tiny() -> Self { + Self { + n_mels: 80, + n_vocab: 51864, + n_audio_ctx: 1500, + n_audio_state: 384, + n_audio_head: 6, + n_audio_layer: 4, + n_text_ctx: 448, + n_text_state: 384, + n_text_head: 6, + n_text_layer: 4, + } + } +} + struct Embedding { embeddings: Tensor, hidden_size: usize, From e6b01d0c18d27b2363f5aad7a19da38afc51f7d1 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 10:01:05 +0100 Subject: [PATCH 08/17] Add the conv1d layer (but not the op). --- candle-examples/examples/whisper/main.rs | 89 ++++++++++++++++++++++-- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index dac6df84..75ab2189 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -172,6 +172,79 @@ impl Linear { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ConvConfig { + padding: usize, + stride: usize, +} + +impl Default for ConvConfig { + fn default() -> Self { + Self { + padding: 0, + stride: 1, + } + } +} + +struct Conv1D { + weight: Tensor, + bias: Option, + config: ConvConfig, +} + +impl Conv1D { + fn load( + in_channels: usize, + out_channels: usize, + kernel_size: usize, + config: ConvConfig, + p: &str, + vb: &VarBuilder, + ) -> Result { + let weight = vb.get( + (out_channels, in_channels, kernel_size), + &format!("{p}.weight"), + )?; + let bias = vb.get(out_channels, &format!("{p}.bias"))?; + Ok(Self { + weight, + bias: Some(bias), + config, + }) + } + + fn load_no_bias( + in_channels: usize, + out_channels: usize, + kernel_size: usize, + config: ConvConfig, + p: &str, + vb: &VarBuilder, + ) -> Result { + let weight = vb.get( + (out_channels, in_channels, kernel_size), + &format!("{p}.weight"), + )?; + Ok(Self { + weight, + bias: None, + config, + }) + } + + fn forward(&self, x: &Tensor) -> candle::Result { + let (bsize, _, _) = x.shape().r3()?; + let w = self.weight.broadcast_left(bsize)?.t()?; + // TODO: Add the conv1d operation + let x = x.matmul(&w)?; + match &self.bias { + None => Ok(x), + Some(bias) => x.broadcast_add(bias), + } + } +} + struct Dropout { pr: f64, } @@ -341,8 +414,8 @@ fn sinusoids(length: usize, channels: usize) -> Result { // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 struct AudioEncoder { - conv1: Linear, // TODO - conv2: Linear, // TODO + conv1: Conv1D, + conv2: Conv1D, positional_embedding: Tensor, blocks: Vec, ln_post: LayerNorm, @@ -352,8 +425,16 @@ impl AudioEncoder { fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { let n_state = cfg.n_audio_state; let n_head = cfg.n_audio_head; - let conv1 = Linear::load(cfg.n_mels, n_state, &format!("{p}.conv1"), vb)?; - let conv2 = Linear::load(n_state, n_state, &format!("{p}.conv2"), vb)?; + let cfg1 = ConvConfig { + padding: 1, + stride: 1, + }; + let cfg2 = ConvConfig { + padding: 1, + stride: 2, + }; + let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?; + let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?; let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?; let blocks = (0..cfg.n_audio_layer) .map(|i| { From 3aac1047fec43a4d756ae4e60a8ae82f7c3e636e Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 10:52:34 +0100 Subject: [PATCH 09/17] Sketch the conv1d op. --- candle-core/src/backprop.rs | 8 +++++++- candle-core/src/cpu_backend.rs | 11 +++++++++++ candle-core/src/cuda_backend.rs | 11 +++++++++++ candle-core/src/dummy_cuda_backend.rs | 11 +++++++++++ candle-core/src/op.rs | 8 ++++++++ candle-core/src/storage.rs | 27 +++++++++++++++++++++++++++ candle-core/src/tensor.rs | 22 ++++++++++++++++++++++ 7 files changed, 97 insertions(+), 1 deletion(-) diff --git a/candle-core/src/backprop.rs b/candle-core/src/backprop.rs index 45448505..a44f732f 100644 --- a/candle-core/src/backprop.rs +++ b/candle-core/src/backprop.rs @@ -33,7 +33,12 @@ impl Tensor { track_grad |= tg; nodes } - Op::Add(lhs, rhs) + Op::Conv1D { + arg: lhs, + kernel: rhs, + .. + } + | Op::Add(lhs, rhs) | Op::Mul(lhs, rhs) | Op::Sub(lhs, rhs) | Op::Div(lhs, rhs) @@ -147,6 +152,7 @@ impl Tensor { let f_grad = pred.where_cond(&zeros, &grad)?; *f_sum_grad = f_sum_grad.add(&f_grad)?; } + Op::Conv1D { .. } => return Err(Error::BackwardNotSupported { op: "conv1d" }), Op::Embedding(_lhs, _rhs) => { return Err(Error::BackwardNotSupported { op: "embedding" }) } diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index 0871175f..ed3a5998 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -627,6 +627,17 @@ impl CpuStorage { WCond(pred, layout).map(t, t_l, f, f_l) } + pub(crate) fn conv1d( + &self, + _l: &Layout, + _kernel: &Self, + _kernel_l: &Layout, + _padding: usize, + _stride: usize, + ) -> Result { + todo!() + } + pub(crate) fn embedding(&self, ids_l: &Layout, rhs: &Self, rhs_l: &Layout) -> Result { let ids = self.as_slice::()?; let (vocab_size, hidden_size) = rhs_l.shape().r2()?; diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs index 0c87004b..ec69688c 100644 --- a/candle-core/src/cuda_backend.rs +++ b/candle-core/src/cuda_backend.rs @@ -801,6 +801,17 @@ impl CudaStorage { Ok(Self { slice, device }) } + pub(crate) fn conv1d( + &self, + _l: &Layout, + _kernel: &Self, + _kernel_l: &Layout, + _padding: usize, + _stride: usize, + ) -> Result { + todo!() + } + pub(crate) fn embedding(&self, layout: &Layout, rhs: &Self, rhs_l: &Layout) -> Result { let device = self.device().clone(); let slice = Embedding(self, layout).map(&rhs.slice, &device, rhs_l)?; diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs index b025eeab..eca5961b 100644 --- a/candle-core/src/dummy_cuda_backend.rs +++ b/candle-core/src/dummy_cuda_backend.rs @@ -100,6 +100,17 @@ impl CudaStorage { Err(Error::NotCompiledWithCudaSupport) } + pub(crate) fn conv1d( + &self, + _l: &Layout, + _kernel: &Self, + _kernel_l: &Layout, + _padding: usize, + _stride: usize, + ) -> Result { + Err(Error::NotCompiledWithCudaSupport) + } + pub(crate) fn embedding(&self, _: &Layout, _: &Self, _: &Layout) -> Result { Err(Error::NotCompiledWithCudaSupport) } diff --git a/candle-core/src/op.rs b/candle-core/src/op.rs index 860be0b3..ee57b325 100644 --- a/candle-core/src/op.rs +++ b/candle-core/src/op.rs @@ -12,6 +12,14 @@ pub(crate) enum Op { Embedding(Tensor, Tensor), WhereCond(Tensor, Tensor, Tensor), + #[allow(dead_code)] + Conv1D { + arg: Tensor, + kernel: Tensor, + padding: usize, + stride: usize, + }, + Cat(Vec, usize), #[allow(dead_code)] // add is currently unused. diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs index 4e630a58..235080c0 100644 --- a/candle-core/src/storage.rs +++ b/candle-core/src/storage.rs @@ -144,6 +144,33 @@ impl Storage { } } + pub(crate) fn conv1d( + &self, + l: &Layout, + kernel: &Self, + kernel_l: &Layout, + padding: usize, + stride: usize, + ) -> Result { + self.same_device(kernel, "conv1d")?; + self.same_dtype(kernel, "conv1d")?; + match (self, &kernel) { + (Storage::Cpu(inp), Storage::Cpu(kernel)) => { + let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + Ok(Self::Cpu(s)) + } + (Storage::Cuda(inp), Storage::Cuda(kernel)) => { + let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + Ok(Self::Cuda(s)) + } + (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp { + lhs: lhs.device().location(), + rhs: rhs.device().location(), + op: "conv1d", + }), + } + } + pub(crate) fn where_cond( &self, layout: &Layout, diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs index a468d879..26d44718 100644 --- a/candle-core/src/tensor.rs +++ b/candle-core/src/tensor.rs @@ -432,6 +432,28 @@ impl Tensor { Ok(from_storage(storage, dims, op, false)) } + pub fn conv1d(&self, kernel: &Self, padding: usize, stride: usize) -> Result { + let storage = self.storage.conv1d( + self.layout(), + &kernel.storage, + kernel.layout(), + padding, + stride, + )?; + let op = if self.track_op() || kernel.track_op() { + Some(Op::Conv1D { + arg: self.clone(), + kernel: kernel.clone(), + padding, + stride, + }) + } else { + None + }; + let dims = self.dims(); + Ok(from_storage(storage, dims, op, false)) + } + pub fn matmul(&self, rhs: &Self) -> Result { let a_dims = self.shape().dims(); let b_dims = rhs.shape().dims(); From a424d95473ea9268ffb1dde4d73ce0cff9904845 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 11:15:45 +0100 Subject: [PATCH 10/17] Add more of the conv1d op. --- candle-core/src/conv.rs | 24 +++++++++++++++++++++ candle-core/src/cpu_backend.rs | 3 +-- candle-core/src/cuda_backend.rs | 3 +-- candle-core/src/dummy_cuda_backend.rs | 3 +-- candle-core/src/lib.rs | 1 + candle-core/src/storage.rs | 7 +++--- candle-core/src/tensor.rs | 27 ++++++++++++++++++------ candle-examples/examples/whisper/main.rs | 3 +-- 8 files changed, 52 insertions(+), 19 deletions(-) create mode 100644 candle-core/src/conv.rs diff --git a/candle-core/src/conv.rs b/candle-core/src/conv.rs new file mode 100644 index 00000000..90bb5229 --- /dev/null +++ b/candle-core/src/conv.rs @@ -0,0 +1,24 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ParamsConv1D { + pub(crate) b_size: Option, + pub(crate) c_out: usize, + pub(crate) c_in: usize, + pub(crate) k_size: usize, + pub(crate) padding: usize, + pub(crate) stride: usize, +} + +impl ParamsConv1D { + pub(crate) fn l_out(&self, l_in: usize) -> usize { + let dilation = 1; + (l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1 + } + + pub(crate) fn out_dims(&self, l_in: usize) -> Vec { + let l_out = self.l_out(l_in); + match self.b_size { + None => vec![self.c_out, l_out], + Some(n) => vec![n, self.c_out, l_out], + } + } +} diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index ed3a5998..54002184 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -632,8 +632,7 @@ impl CpuStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { todo!() } diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs index ec69688c..917655fc 100644 --- a/candle-core/src/cuda_backend.rs +++ b/candle-core/src/cuda_backend.rs @@ -806,8 +806,7 @@ impl CudaStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { todo!() } diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs index eca5961b..0dbd8d54 100644 --- a/candle-core/src/dummy_cuda_backend.rs +++ b/candle-core/src/dummy_cuda_backend.rs @@ -105,8 +105,7 @@ impl CudaStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { Err(Error::NotCompiledWithCudaSupport) } diff --git a/candle-core/src/lib.rs b/candle-core/src/lib.rs index 0d4c2a8d..2365a34d 100644 --- a/candle-core/src/lib.rs +++ b/candle-core/src/lib.rs @@ -1,4 +1,5 @@ mod backprop; +mod conv; mod cpu_backend; #[cfg(feature = "cuda")] mod cuda_backend; diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs index 235080c0..53ea1544 100644 --- a/candle-core/src/storage.rs +++ b/candle-core/src/storage.rs @@ -149,18 +149,17 @@ impl Storage { l: &Layout, kernel: &Self, kernel_l: &Layout, - padding: usize, - stride: usize, + params: &crate::conv::ParamsConv1D, ) -> Result { self.same_device(kernel, "conv1d")?; self.same_dtype(kernel, "conv1d")?; match (self, &kernel) { (Storage::Cpu(inp), Storage::Cpu(kernel)) => { - let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + let s = inp.conv1d(l, kernel, kernel_l, params)?; Ok(Self::Cpu(s)) } (Storage::Cuda(inp), Storage::Cuda(kernel)) => { - let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + let s = inp.conv1d(l, kernel, kernel_l, params)?; Ok(Self::Cuda(s)) } (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp { diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs index 26d44718..590b81c4 100644 --- a/candle-core/src/tensor.rs +++ b/candle-core/src/tensor.rs @@ -433,13 +433,26 @@ impl Tensor { } pub fn conv1d(&self, kernel: &Self, padding: usize, stride: usize) -> Result { - let storage = self.storage.conv1d( - self.layout(), - &kernel.storage, - kernel.layout(), + let (c_out, c_in_k, k_size) = kernel.shape().r3()?; + let (b_size, c_in, l_in) = match *self.dims() { + [b_size, c_in, l_in] => (Some(b_size), c_in, l_in), + [c_in, l_in] => (None, c_in, l_in), + _ => todo!("proper error message"), + }; + if c_in != c_in_k { + todo!("proper error message") + } + let params = crate::conv::ParamsConv1D { + b_size, + c_out, + c_in, + k_size, padding, stride, - )?; + }; + let storage = + self.storage + .conv1d(self.layout(), &kernel.storage, kernel.layout(), ¶ms)?; let op = if self.track_op() || kernel.track_op() { Some(Op::Conv1D { arg: self.clone(), @@ -450,8 +463,8 @@ impl Tensor { } else { None }; - let dims = self.dims(); - Ok(from_storage(storage, dims, op, false)) + let out_dims = params.out_dims(l_in); + Ok(from_storage(storage, out_dims, op, false)) } pub fn matmul(&self, rhs: &Self) -> Result { diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 75ab2189..a380d30e 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -236,8 +236,7 @@ impl Conv1D { fn forward(&self, x: &Tensor) -> candle::Result { let (bsize, _, _) = x.shape().r3()?; let w = self.weight.broadcast_left(bsize)?.t()?; - // TODO: Add the conv1d operation - let x = x.matmul(&w)?; + let x = x.conv1d(&w, self.config.padding, self.config.stride)?; match &self.bias { None => Ok(x), Some(bias) => x.broadcast_add(bias), From 950b4af49e56b640b87eb273e839b2fd466e1424 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 11:29:28 +0100 Subject: [PATCH 11/17] Proper conv1d dispatch. --- candle-core/src/conv.rs | 11 +++++++---- candle-core/src/cpu_backend.rs | 30 +++++++++++++++++++++++++----- candle-core/src/tensor.rs | 3 ++- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/candle-core/src/conv.rs b/candle-core/src/conv.rs index 90bb5229..041bb6fb 100644 --- a/candle-core/src/conv.rs +++ b/candle-core/src/conv.rs @@ -1,6 +1,9 @@ #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct ParamsConv1D { pub(crate) b_size: Option, + // Maybe we should have a version without l_in as this bit depends on the input and not only on + // the weights. + pub(crate) l_in: usize, pub(crate) c_out: usize, pub(crate) c_in: usize, pub(crate) k_size: usize, @@ -9,13 +12,13 @@ pub(crate) struct ParamsConv1D { } impl ParamsConv1D { - pub(crate) fn l_out(&self, l_in: usize) -> usize { + pub(crate) fn l_out(&self) -> usize { let dilation = 1; - (l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1 + (self.l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1 } - pub(crate) fn out_dims(&self, l_in: usize) -> Vec { - let l_out = self.l_out(l_in); + pub(crate) fn out_dims(&self) -> Vec { + let l_out = self.l_out(); match self.b_size { None => vec![self.c_out, l_out], Some(n) => vec![n, self.c_out, l_out], diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index 54002184..718b071c 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -202,6 +202,26 @@ fn copy_strided_src_( } } +struct Conv1D<'a>(&'a crate::conv::ParamsConv1D); + +impl<'a> Map2 for Conv1D<'a> { + const OP: &'static str = "conv1d"; + fn f( + &self, + _inp: &[T], + _inp_l: &Layout, + _k: &[T], + _k_l: &Layout, + ) -> Result> { + let p = self.0; + let l_out = p.l_out(); + let out_elems = p.c_out * l_out * p.b_size.unwrap_or(1); + let dst = vec![T::zero(); out_elems]; + // TODO: actually implement the ops. + Ok(dst) + } +} + struct MatMul((usize, usize, usize, usize)); impl Map2 for MatMul { @@ -629,12 +649,12 @@ impl CpuStorage { pub(crate) fn conv1d( &self, - _l: &Layout, - _kernel: &Self, - _kernel_l: &Layout, - _params: &crate::conv::ParamsConv1D, + l: &Layout, + kernel: &Self, + kernel_l: &Layout, + params: &crate::conv::ParamsConv1D, ) -> Result { - todo!() + Conv1D(params).map(self, l, kernel, kernel_l) } pub(crate) fn embedding(&self, ids_l: &Layout, rhs: &Self, rhs_l: &Layout) -> Result { diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs index 590b81c4..25ab0a9b 100644 --- a/candle-core/src/tensor.rs +++ b/candle-core/src/tensor.rs @@ -444,6 +444,7 @@ impl Tensor { } let params = crate::conv::ParamsConv1D { b_size, + l_in, c_out, c_in, k_size, @@ -463,7 +464,7 @@ impl Tensor { } else { None }; - let out_dims = params.out_dims(l_in); + let out_dims = params.out_dims(); Ok(from_storage(storage, out_dims, op, false)) } From aea090401ded5789e95f1f8efb7404a66b508356 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 12:03:28 +0100 Subject: [PATCH 12/17] Run the text decoding bit. --- .../examples/whisper/extract_weights.py | 2 +- candle-examples/examples/whisper/main.rs | 39 +++++++++++-------- 2 files changed, 24 insertions(+), 17 deletions(-) diff --git a/candle-examples/examples/whisper/extract_weights.py b/candle-examples/examples/whisper/extract_weights.py index d6ccffc6..65602703 100644 --- a/candle-examples/examples/whisper/extract_weights.py +++ b/candle-examples/examples/whisper/extract_weights.py @@ -8,6 +8,6 @@ data = torch.load("tiny.en.pt") weights = {} for k, v in data["model_state_dict"].items(): weights[k] = v.contiguous() - print(k, v.shape) + print(k, v.shape, v.dtype) save_file(weights, "tiny.en.safetensors") print(data["dims"]) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index a380d30e..c3a2769f 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -3,7 +3,7 @@ // TODO: // - kv-cache support? -use anyhow::{Error as E, Result}; +use anyhow::Result; use candle::{safetensors::SafeTensors, DType, Device, Shape, Tensor}; use clap::Parser; use std::collections::HashMap; @@ -97,7 +97,7 @@ struct Config { } impl Config { - fn tiny() -> Self { + fn tiny_en() -> Self { Self { n_mels: 80, n_vocab: 51864, @@ -302,8 +302,8 @@ struct MultiHeadAttention { impl MultiHeadAttention { fn load(n_state: usize, n_head: usize, p: &str, vb: &VarBuilder) -> Result { let query = Linear::load(n_state, n_state, &format!("{p}.query"), vb)?; - let value = Linear::load_no_bias(n_state, n_state, &format!("{p}.value"), vb)?; - let key = Linear::load(n_state, n_state, &format!("{p}.key"), vb)?; + let value = Linear::load(n_state, n_state, &format!("{p}.value"), vb)?; + let key = Linear::load_no_bias(n_state, n_state, &format!("{p}.key"), vb)?; let out = Linear::load(n_state, n_state, &format!("{p}.out"), vb)?; Ok(Self { query, @@ -500,12 +500,13 @@ impl TextDecoder { let last = x_dims[x_dims.len() - 1]; let token_embedding = self.token_embedding.forward(x)?; let positional_embedding = self.positional_embedding.narrow(0, 0, last)?; - let mut x = (token_embedding + positional_embedding)?; + let mut x = token_embedding.broadcast_add(&positional_embedding)?; for block in self.blocks.iter() { x = block.forward(&x, Some(xa))? } let x = self.ln.forward(&x)?; - let logits = x.matmul(&self.token_embedding.embeddings.t()?)?; + let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?; + let logits = x.matmul(&w.t()?)?; Ok(logits) } } @@ -517,9 +518,9 @@ struct Whisper { } impl Whisper { - fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { - let encoder = AudioEncoder::load(&format!("{p}.encoder"), vb, cfg)?; - let decoder = TextDecoder::load(&format!("{p}.decoder"), vb, cfg)?; + fn load(vb: &VarBuilder, cfg: &Config) -> Result { + let encoder = AudioEncoder::load("encoder", vb, cfg)?; + let decoder = TextDecoder::load("decoder", vb, cfg)?; Ok(Self { encoder, decoder }) } fn forward(&self, mel: &Tensor, tokens: &Tensor) -> Result { @@ -537,15 +538,13 @@ struct Args { cpu: bool, #[arg(long)] - tokenizer_config: String, + weights: String, #[arg(long)] - weights: String, + input: String, } fn main() -> Result<()> { - use tokenizers::Tokenizer; - let args = Args::parse(); let device = if args.cpu { Device::Cpu @@ -553,11 +552,19 @@ fn main() -> Result<()> { Device::new_cuda(0)? }; - let mut tokenizer = Tokenizer::from_file(args.tokenizer_config).map_err(E::msg)?; - let _tokenizer = tokenizer.with_padding(None).with_truncation(None); + let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? }; + let input = input.deserialize()?; + let x = input.tensor("x", &device)?.to_dtype(DType::U32)?; + let xa = input.tensor("xa", &device)?; let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? }; let weights = weights.deserialize()?; - let _vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device); + let vb = VarBuilder::from_safetensors(vec![weights], DTYPE, device.clone()); + let cfg = Config::tiny_en(); + + let model = Whisper::load(&vb, &cfg)?; + let logits = model.decoder.forward(&x, &xa)?; + println!("{logits}"); + println!("python logits: {}", input.tensor("logits", &device)?); Ok(()) } From 0d99b4379224838265bf9bedd836dd655b94a001 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 12:27:05 +0100 Subject: [PATCH 13/17] Line up the textdecoder values with the python implementation. --- candle-examples/examples/whisper/main.rs | 36 ++++++++++++------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index c3a2769f..6341c5ee 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -323,13 +323,18 @@ impl MultiHeadAttention { Ok(out) } - fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { - let (n_batch, n_ctx, n_state) = q.shape().r3()?; + fn reshape_head(&self, x: &Tensor) -> Result { + let (n_batch, n_ctx, n_state) = x.shape().r3()?; let target_dims = &[n_batch, n_ctx, self.n_head, n_state / self.n_head]; + Ok(x.reshape(target_dims)?.transpose(1, 2)?) + } + + fn qkv_attention(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let (_, _, n_state) = q.shape().r3()?; let scale = ((n_state / self.n_head) as f64).powf(-0.25); - let q = (q.reshape(target_dims)?.transpose(1, 2)? * scale)?; - let k = (k.reshape(target_dims)?.transpose(1, 2)?.transpose(2, 3)? * scale)?; - let v = v.reshape(target_dims)?.transpose(1, 2)?; + let q = (self.reshape_head(q)? * scale)?; + let k = (self.reshape_head(k)?.transpose(2, 3)? * scale)?; + let v = self.reshape_head(v)?.contiguous()?; let qk = q.matmul(&k)?; let w = qk.softmax(qk.rank() - 1)?; let wv = w.matmul(&v)?.transpose(1, 2)?.flatten(Some(2), None)?; @@ -341,8 +346,7 @@ impl MultiHeadAttention { struct ResidualAttentionBlock { attn: MultiHeadAttention, attn_ln: LayerNorm, - cross_attn: Option, - cross_attn_ln: Option, + cross_attn: Option<(MultiHeadAttention, LayerNorm)>, mlp_linear1: Linear, mlp_linear2: Linear, mlp_ln: LayerNorm, @@ -352,13 +356,13 @@ impl ResidualAttentionBlock { fn load(n_state: usize, n_head: usize, ca: bool, p: &str, vb: &VarBuilder) -> Result { let attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.attn"), vb)?; let attn_ln = LayerNorm::load(n_state, &format!("{p}.attn_ln"), vb)?; - let (cross_attn, cross_attn_ln) = if ca { + let cross_attn = if ca { let cross_attn = MultiHeadAttention::load(n_state, n_head, &format!("{p}.cross_attn"), vb)?; let cross_attn_ln = LayerNorm::load(n_state, &format!("{p}.cross_attn_ln"), vb)?; - (Some(cross_attn), Some(cross_attn_ln)) + Some((cross_attn, cross_attn_ln)) } else { - (None, None) + None }; let n_mlp = n_state * 4; let mlp_linear1 = Linear::load(n_state, n_mlp, &format!("{p}.mlp.0"), vb)?; @@ -368,7 +372,6 @@ impl ResidualAttentionBlock { attn, attn_ln, cross_attn, - cross_attn_ln, mlp_linear1, mlp_linear2, mlp_ln, @@ -379,11 +382,8 @@ impl ResidualAttentionBlock { let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?; let mut x = (x + attn)?; // Cross-Attn - if let Some(cross_attn_ln) = &self.cross_attn_ln { - x = cross_attn_ln.forward(&x)? - } - if let Some(cross_attn) = &self.cross_attn { - x = cross_attn.forward(&x, xa)? + if let Some((attn, ln)) = &self.cross_attn { + x = (&x + attn.forward(&ln.forward(&x)?, xa)?)?; } // Mlp let mlp = self.mlp_linear2.forward( @@ -482,7 +482,7 @@ impl TextDecoder { vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))?; let blocks = (0..cfg.n_text_layer) .map(|i| { - ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) + ResidualAttentionBlock::load(n_state, n_head, true, &format!("{p}.blocks.{i}"), vb) }) .collect::>>()?; let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?; @@ -502,7 +502,7 @@ impl TextDecoder { let positional_embedding = self.positional_embedding.narrow(0, 0, last)?; let mut x = token_embedding.broadcast_add(&positional_embedding)?; for block in self.blocks.iter() { - x = block.forward(&x, Some(xa))? + x = block.forward(&x, Some(xa))?; } let x = self.ln.forward(&x)?; let w = self.token_embedding.embeddings.broadcast_left(x_dims[0])?; From 599160605c0294c94c33f64aeca0ac9f388d03c7 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 13:13:28 +0100 Subject: [PATCH 14/17] Use the stored embeddings. --- candle-examples/examples/whisper/main.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 6341c5ee..1b6f4bfe 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -424,6 +424,7 @@ impl AudioEncoder { fn load(p: &str, vb: &VarBuilder, cfg: &Config) -> Result { let n_state = cfg.n_audio_state; let n_head = cfg.n_audio_head; + let n_ctx = cfg.n_audio_ctx; let cfg1 = ConvConfig { padding: 1, stride: 1, @@ -434,7 +435,12 @@ impl AudioEncoder { }; let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?; let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?; - let positional_embedding = sinusoids(cfg.n_audio_ctx, n_state)?.to_device(&vb.device)?; + /* The positional embeddings could be regenerated via the following. */ + let positional_embedding = if true { + vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))? + } else { + sinusoids(n_ctx, n_state)?.to_device(&vb.device)? + }; let blocks = (0..cfg.n_audio_layer) .map(|i| { ResidualAttentionBlock::load(n_state, n_head, false, &format!("{p}.blocks.{i}"), vb) From 99b83773b51382915c94bd3b2298522036235f72 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 13:21:59 +0100 Subject: [PATCH 15/17] Small cleanup. --- candle-examples/examples/whisper/main.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 1b6f4bfe..71b03e72 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -435,10 +435,10 @@ impl AudioEncoder { }; let conv1 = Conv1D::load(cfg.n_mels, n_state, 3, cfg1, &format!("{p}.conv1"), vb)?; let conv2 = Conv1D::load(n_state, n_state, 3, cfg2, &format!("{p}.conv2"), vb)?; - /* The positional embeddings could be regenerated via the following. */ let positional_embedding = if true { vb.get((n_ctx, n_state), &format!("{p}.positional_embedding"))? } else { + /* The positional embeddings could be regenerated via the following. */ sinusoids(n_ctx, n_state)?.to_device(&vb.device)? }; let blocks = (0..cfg.n_audio_layer) @@ -474,7 +474,6 @@ struct TextDecoder { positional_embedding: Tensor, blocks: Vec, ln: LayerNorm, - mask: Tensor, } impl TextDecoder { @@ -492,13 +491,11 @@ impl TextDecoder { }) .collect::>>()?; let ln = LayerNorm::load(n_state, &format!("{p}.ln"), vb)?; - let mask = Tensor::new(&[0u32], &vb.device)?; // TODO Ok(Self { token_embedding, positional_embedding, blocks, ln, - mask, }) } fn forward(&self, x: &Tensor, xa: &Tensor) -> Result { From b3d4d0fd0f7c0115cf5c8cea60094abef5536e56 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 13:50:41 +0100 Subject: [PATCH 16/17] Very inefficient conv1d implementation. --- candle-core/src/cpu_backend.rs | 53 +++++++++++++++++++++++++++++----- 1 file changed, 45 insertions(+), 8 deletions(-) diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index 718b071c..4eb57bc7 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -206,18 +206,55 @@ struct Conv1D<'a>(&'a crate::conv::ParamsConv1D); impl<'a> Map2 for Conv1D<'a> { const OP: &'static str = "conv1d"; - fn f( + fn f( &self, - _inp: &[T], - _inp_l: &Layout, - _k: &[T], - _k_l: &Layout, + inp: &[T], + inp_l: &Layout, + k: &[T], + k_l: &Layout, ) -> Result> { + // TODO: Optimize this (proper algorithm, simd, multithread, remove bound checks, etc). let p = self.0; + let inp = &inp[inp_l.start_offset()..]; + let k = &k[k_l.start_offset()..]; + let inp_stride = inp_l.stride(); + let (inp_stride0, inp_stride) = if inp_stride.len() == 3 { + (inp_stride[0], &inp_stride[1..]) + } else { + (0, inp_stride) // This value never gets used anyway + }; + let k_stride = k_l.stride(); + let k_over_2 = p.k_size / 2; let l_out = p.l_out(); - let out_elems = p.c_out * l_out * p.b_size.unwrap_or(1); - let dst = vec![T::zero(); out_elems]; - // TODO: actually implement the ops. + let dst_elems = p.c_out * l_out * p.b_size.unwrap_or(1); + let mut dst = vec![T::zero(); dst_elems]; + // The output shape is [b_size, c_out, l_out] + for b_idx in 0..p.b_size.unwrap_or(1) { + let inp_idx = b_idx * inp_stride0; + let dst_idx = b_idx * p.c_out * l_out; + for dst_c_idx in 0..p.c_out { + let dst_idx = dst_idx + dst_c_idx * l_out; + for dst_l in 0..l_out { + let dst_idx = dst_idx + dst_l; + let mut d = T::zero(); + for offset in 0..p.k_size { + // inp[bidx, src_c_idx, dst_l + offset - k//2] * k[dst_c_idx, src_c_idx, offset] + if k_over_2 <= dst_l + offset && dst_l + offset < k_over_2 + p.l_in { + let src_l = dst_l + offset - k_over_2; + for src_c_idx in 0..p.c_in { + let inp_idx = + inp_idx + src_c_idx * inp_stride[0] + src_l * inp_stride[1]; + let k_idx = dst_c_idx * k_stride[0] + + src_c_idx * k_stride[1] + + offset * k_stride[2]; + d += inp[inp_idx] * k[k_idx] + } + } + } + dst[dst_idx] = d + } + } + } Ok(dst) } } From c3739d001bfb1a1305fc1ff398761e380f87bfb1 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 14:06:09 +0100 Subject: [PATCH 17/17] Get the audio-encoder to return some values. --- candle-examples/examples/whisper/main.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 71b03e72..839dfc13 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -233,13 +233,15 @@ impl Conv1D { }) } - fn forward(&self, x: &Tensor) -> candle::Result { - let (bsize, _, _) = x.shape().r3()?; - let w = self.weight.broadcast_left(bsize)?.t()?; - let x = x.conv1d(&w, self.config.padding, self.config.stride)?; + fn forward(&self, x: &Tensor) -> Result { + let x = x.conv1d(&self.weight, self.config.padding, self.config.stride)?; match &self.bias { None => Ok(x), - Some(bias) => x.broadcast_add(bias), + Some(bias) => { + let b = bias.shape().r1()?; + let bias = bias.reshape((1, b, 1))?; + Ok(x.broadcast_add(&bias)?) + } } } } @@ -381,11 +383,9 @@ impl ResidualAttentionBlock { fn forward(&self, x: &Tensor, xa: Option<&Tensor>) -> Result { let attn = self.attn.forward(&self.attn_ln.forward(x)?, None)?; let mut x = (x + attn)?; - // Cross-Attn if let Some((attn, ln)) = &self.cross_attn { x = (&x + attn.forward(&ln.forward(&x)?, xa)?)?; } - // Mlp let mlp = self.mlp_linear2.forward( &self .mlp_linear1 @@ -557,8 +557,8 @@ fn main() -> Result<()> { let input = unsafe { candle::safetensors::MmapedFile::new(args.input)? }; let input = input.deserialize()?; - let x = input.tensor("x", &device)?.to_dtype(DType::U32)?; - let xa = input.tensor("xa", &device)?; + let tokens = input.tensor("tokens", &device)?.to_dtype(DType::U32)?; + let mel = input.tensor("mel", &device)?; let weights = unsafe { candle::safetensors::MmapedFile::new(args.weights)? }; let weights = weights.deserialize()?; @@ -566,8 +566,8 @@ fn main() -> Result<()> { let cfg = Config::tiny_en(); let model = Whisper::load(&vb, &cfg)?; - let logits = model.decoder.forward(&x, &xa)?; + let logits = model.forward(&mel, &tokens)?; println!("{logits}"); - println!("python logits: {}", input.tensor("logits", &device)?); + println!("python logits: {}", input.tensor("dec", &device)?); Ok(()) }