-
Notifications
You must be signed in to change notification settings - Fork 211
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Implement GTE model to support the non-flash-attn version (#446)
Co-authored-by: Hyeongchan Kim <[email protected]>
- Loading branch information
1 parent
e27a4fb
commit 0462171
Showing
13 changed files
with
3,964 additions
and
203 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| use candle::{DType, Device, Result, Tensor, D}; | ||
| use serde::Deserialize; | ||
|
|
||
| #[derive(Debug, Clone, PartialEq, Deserialize)] | ||
| pub struct NTKScaling { | ||
| pub factor: f32, | ||
| } | ||
|
|
||
| #[derive(Debug, Clone, PartialEq, Deserialize)] | ||
| #[serde(tag = "type", rename_all = "kebab-case")] | ||
| pub enum RopeScaling { | ||
| Ntk(NTKScaling), | ||
| } | ||
|
|
||
| pub fn get_inv_freqs( | ||
| dim: usize, | ||
| base: f32, | ||
| device: &Device, | ||
| rope_scaling: Option<&RopeScaling>, | ||
| ) -> Result<Tensor> { | ||
| let get_inv_freqs_inner = |dim: usize, base: f32, device: &Device| { | ||
| let inv_freq: Vec<_> = (0..dim) | ||
| .step_by(2) | ||
| .map(|i| 1f32 / base.powf(i as f32 / dim as f32)) | ||
| .collect(); | ||
| let inv_freq_len = inv_freq.len(); | ||
| Tensor::from_vec(inv_freq, (1, inv_freq_len), device) | ||
| }; | ||
|
|
||
| if let Some(rope_scaling) = rope_scaling { | ||
| match rope_scaling { | ||
| RopeScaling::Ntk(ntk_scaling) => { | ||
| let inv_freqs = get_inv_freqs_inner(dim, base * ntk_scaling.factor, device)?; | ||
| let s = ntk_scaling.factor.powf(2.0 / dim as f32) as f64; | ||
| return inv_freqs / s; | ||
| } | ||
| } | ||
| } | ||
| get_inv_freqs_inner(dim, base, device) | ||
| } | ||
|
|
||
| pub fn get_cos_sin( | ||
| length: usize, | ||
| inv_freqs: &Tensor, | ||
| dtype: DType, | ||
| repeat_freqs: bool, | ||
| ) -> Result<(Tensor, Tensor)> { | ||
| let t = Tensor::arange(0u32, length as u32, inv_freqs.device())? | ||
| .to_dtype(DType::F32)? | ||
| .reshape((length, 1))?; | ||
| let mut freqs = t.matmul(inv_freqs)?; | ||
| if repeat_freqs { | ||
| freqs = Tensor::cat(&[&freqs, &freqs], 1)?; | ||
| } | ||
|
|
||
| let cos = freqs.cos()?.to_dtype(dtype)?; | ||
| let sin = freqs.sin()?.to_dtype(dtype)?; | ||
| Ok((cos, sin)) | ||
| } | ||
|
|
||
| pub fn apply_rotary( | ||
| x: &Tensor, | ||
| cos: &Tensor, | ||
| sin: &Tensor, | ||
| attention_head_size: usize, | ||
| ) -> Result<Tensor> { | ||
| let dim = attention_head_size / 2; | ||
| let x1 = x.narrow(D::Minus1, 0, dim)?; | ||
| let x2 = x.narrow(D::Minus1, dim, dim)?; | ||
| let rotate_x = Tensor::cat(&[&x2.neg()?, &x1], D::Minus1)?; | ||
| let rope = (x.broadcast_mul(cos)? + rotate_x.broadcast_mul(sin)?)?; | ||
| Ok(rope) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.