diff options
Diffstat (limited to 'src/model.rs')
| -rw-r--r-- | src/model.rs | 333 |
1 files changed, 333 insertions, 0 deletions
diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000..acac139 --- /dev/null +++ b/src/model.rs @@ -0,0 +1,333 @@ +use burn::module::Module; +use burn::module::ParamId; +use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; +use burn::prelude::*; +use burn::tensor::activation::{silu, softmax}; +use burn::tensor::Bool; + +use crate::config::{HfConfig, LoraConfig}; +use crate::lora::LoraLinear; + +/// RoPE with the Hugging Face Llama "half-split" convention: head_dim is +/// split in half and rotated as (x1, x2) -> (-x2, x1). +/// Note: `burn::nn::RotaryEncoding` uses the interleaved convention (it +/// rotates adjacent pairs), which is incompatible with these weights. +#[derive(Module, Debug)] +pub struct Rope<B: Backend> { + cos: Tensor<B, 2>, + sin: Tensor<B, 2>, +} + +impl<B: Backend> Rope<B> { + pub fn new(head_dim: usize, max_seq: usize, theta: f32, device: &B::Device) -> Self { + let half = head_dim / 2; + let inv_freq: Vec<f32> = (0..half) + .map(|i| 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32)) + .collect(); + let inv_freq = Tensor::<B, 1>::from_floats(inv_freq.as_slice(), device); + let pos = Tensor::<B, 1, Int>::arange(0..max_seq as i64, device).float(); + + // [max_seq, half] = outer(pos, inv_freq) + let freqs = pos + .unsqueeze_dim::<2>(1) + .matmul(inv_freq.unsqueeze_dim::<2>(0)); + // duplicate to cover the whole head_dim: [f0..fh, f0..fh] + let emb = Tensor::cat(vec![freqs.clone(), freqs], 1); + + Self { + cos: emb.clone().cos(), + sin: emb.sin(), + } + } + + /// x: [batch, heads, seq, head_dim] + pub fn apply(&self, x: Tensor<B, 4>) -> Tensor<B, 4> { + let [b, h, t, hd] = x.dims(); + let half = hd / 2; + + let cos = self.cos.clone().slice([0..t, 0..hd]).unsqueeze::<4>(); + let sin = self.sin.clone().slice([0..t, 0..hd]).unsqueeze::<4>(); + + let x1 = x.clone().slice([0..b, 0..h, 0..t, 0..half]); + let x2 = x.clone().slice([0..b, 0..h, 0..t, half..hd]); + let rotated = Tensor::cat(vec![-x2, x1], 3); + + x * cos + rotated * sin + } +} + +#[derive(Module, Debug)] +pub struct Attention<B: Backend> { + pub q: LoraLinear<B>, + pub k: LoraLinear<B>, + pub v: LoraLinear<B>, + pub o: LoraLinear<B>, + pub n_heads: usize, + pub n_kv_heads: usize, + pub head_dim: usize, +} + +impl<B: Backend> Attention<B> { + fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self { + let hd = cfg.head_dim(); + let d = cfg.hidden_size; + let kv = cfg.num_key_value_heads * hd; + let r = |name: &str| { + if lora.targets(name) { + Some((lora.rank, lora.scaling())) + } else { + None + } + }; + Self { + q: LoraLinear::new(d, d, r("q"), device), + k: LoraLinear::new(d, kv, r("k"), device), + v: LoraLinear::new(d, kv, r("v"), device), + o: LoraLinear::new(d, d, r("o"), device), + n_heads: cfg.num_attention_heads, + n_kv_heads: cfg.num_key_value_heads, + head_dim: hd, + } + } + + /// Repeats each KV head `groups` times, interleaved, like Hugging Face + /// `repeat_kv`: [k0,k0,k1,k1,...]. + fn repeat_kv(&self, x: Tensor<B, 4>, groups: usize) -> Tensor<B, 4> { + if groups == 1 { + return x; + } + let [b, n_kv, t, hd] = x.dims(); + x.reshape([b, n_kv, 1, t, hd]) + .repeat_dim(2, groups) + .reshape([b, n_kv * groups, t, hd]) + } + + pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> { + let [b, t, d] = x.dims(); + let hd = self.head_dim; + let groups = self.n_heads / self.n_kv_heads; + let device = x.device(); + + let q = self + .q + .forward(x.clone()) + .reshape([b, t, self.n_heads, hd]) + .swap_dims(1, 2); + let k = self + .k + .forward(x.clone()) + .reshape([b, t, self.n_kv_heads, hd]) + .swap_dims(1, 2); + let v = self + .v + .forward(x) + .reshape([b, t, self.n_kv_heads, hd]) + .swap_dims(1, 2); + + let q = rope.apply(q); + let k = rope.apply(k); + + let k = self.repeat_kv(k, groups); + let v = self.repeat_kv(v, groups); + + let scale = 1.0 / (hd as f64).sqrt(); + let scores = q.matmul(k.swap_dims(2, 3)) * scale; + + // Causal mask. Mind burn's semantics: `tril_mask` is `false` on the + // triangle that is KEPT (lower + diagonal) and `true` elsewhere, + // which is what gets filled with -inf. + let mask = Tensor::<B, 2, Bool>::tril_mask([t, t], 0, &device).unsqueeze::<4>(); + let scores = scores.mask_fill(mask, f32::NEG_INFINITY); + let attn = softmax(scores, 3); + + let out = attn.matmul(v).swap_dims(1, 2).reshape([b, t, d]); + self.o.forward(out) + } + + fn freeze_base(self) -> Self { + Self { + q: self.q.freeze_base(), + k: self.k.freeze_base(), + v: self.v.freeze_base(), + o: self.o.freeze_base(), + ..self + } + } + + fn trainable_ids(&self, out: &mut Vec<ParamId>) { + for l in [&self.q, &self.k, &self.v, &self.o] { + l.trainable_ids(out); + } + } +} + +#[derive(Module, Debug)] +pub struct Mlp<B: Backend> { + pub gate: LoraLinear<B>, + pub up: LoraLinear<B>, + pub down: LoraLinear<B>, +} + +impl<B: Backend> Mlp<B> { + fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self { + let (d, i) = (cfg.hidden_size, cfg.intermediate_size); + let r = |name: &str| { + if lora.targets(name) { + Some((lora.rank, lora.scaling())) + } else { + None + } + }; + Self { + gate: LoraLinear::new(d, i, r("gate"), device), + up: LoraLinear::new(d, i, r("up"), device), + down: LoraLinear::new(i, d, r("down"), device), + } + } + + pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> { + let gate = silu(self.gate.forward(x.clone())); + self.down.forward(gate * self.up.forward(x)) + } + + fn freeze_base(self) -> Self { + Self { + gate: self.gate.freeze_base(), + up: self.up.freeze_base(), + down: self.down.freeze_base(), + } + } + + fn trainable_ids(&self, out: &mut Vec<ParamId>) { + for l in [&self.gate, &self.up, &self.down] { + l.trainable_ids(out); + } + } +} + +#[derive(Module, Debug)] +pub struct Block<B: Backend> { + pub attn_norm: RmsNorm<B>, + pub attn: Attention<B>, + pub ffn_norm: RmsNorm<B>, + pub ffn: Mlp<B>, +} + +impl<B: Backend> Block<B> { + fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self { + let norm = || { + RmsNormConfig::new(cfg.hidden_size) + .with_epsilon(cfg.rms_norm_eps) + .init(device) + }; + Self { + attn_norm: norm(), + attn: Attention::new(cfg, lora, device), + ffn_norm: norm(), + ffn: Mlp::new(cfg, lora, device), + } + } + + pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> { + let h = x.clone() + self.attn.forward(self.attn_norm.forward(x), rope); + h.clone() + self.ffn.forward(self.ffn_norm.forward(h)) + } + + fn freeze_base(self) -> Self { + Self { + attn: self.attn.freeze_base(), + ffn: self.ffn.freeze_base(), + ..self + } + } + + fn trainable_ids(&self, out: &mut Vec<ParamId>) { + self.attn.trainable_ids(out); + self.ffn.trainable_ids(out); + } +} + +#[derive(Module, Debug)] +pub struct Llama<B: Backend> { + pub embed: Embedding<B>, + pub layers: Vec<Block<B>>, + pub norm: RmsNorm<B>, + pub rope: Rope<B>, +} + +impl<B: Backend> Llama<B> { + pub fn new(cfg: &HfConfig, lora: &LoraConfig, max_seq: usize, device: &B::Device) -> Self { + Self { + embed: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), + layers: (0..cfg.num_hidden_layers) + .map(|_| Block::new(cfg, lora, device)) + .collect(), + norm: RmsNormConfig::new(cfg.hidden_size) + .with_epsilon(cfg.rms_norm_eps) + .init(device), + rope: Rope::new(cfg.head_dim(), max_seq, cfg.rope_theta, device), + } + } + + /// tokens: [batch, seq] -> logits [batch, seq, vocab] + pub fn forward(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3> { + let mut h = self.embed.forward(tokens); + for layer in &self.layers { + h = layer.forward(h, &self.rope); + } + let h = self.norm.forward(h); + // tied embeddings: lm_head = embed.weight^T + h.matmul(self.embed.weight.val().transpose().unsqueeze()) + } + + /// Freezes all base weights; only the LoRA A/B stay trainable. + pub fn freeze_base(self) -> Self { + let Self { + embed, + layers, + norm, + rope, + } = self; + Self { + embed: Embedding { + weight: embed.weight.set_require_grad(false), + }, + layers: layers.into_iter().map(|l| l.freeze_base()).collect(), + norm: RmsNorm { + gamma: norm.gamma.set_require_grad(false), + epsilon: norm.epsilon, + }, + rope, + } + } + + /// L2 norm of all B matrices together. They start at zero, so this is a + /// fingerprint to check that an adapter was trained and that it loads + /// back identical. + pub fn lora_norm(&self) -> f32 { + let mut acc = 0.0f32; + for l in &self.layers { + let mut push = |lin: &LoraLinear<B>| { + if let Some(ad) = &lin.lora { + let v: f32 = ad.b.val().powf_scalar(2.0).sum().into_scalar().elem(); + acc += v; + } + }; + for lin in [&l.attn.q, &l.attn.k, &l.attn.v, &l.attn.o] { + push(lin); + } + for lin in [&l.ffn.gate, &l.ffn.up, &l.ffn.down] { + push(lin); + } + } + acc.sqrt() + } + + pub fn trainable_ids(&self) -> Vec<ParamId> { + let mut out = Vec::new(); + for l in &self.layers { + l.trainable_ids(&mut out); + } + out + } +} |