chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -129,7 +129,6 @@ impl<B: Backend> Attention<B> {
|
||||
let q = rope.apply(q, offset);
|
||||
let k = rope.apply(k, offset);
|
||||
|
||||
// 合并 KV Cache
|
||||
let (k_full, v_full) = match cache {
|
||||
Some(cache) => {
|
||||
let k_full = Tensor::cat(vec![cache.k.clone(), k], 2);
|
||||
|
||||
@@ -4,15 +4,15 @@ use burn::tensor::{Float, Tensor};
|
||||
|
||||
use super::attention::{Attention, KVCache};
|
||||
use super::ffn::FeedForward;
|
||||
use super::norm::RmsNorm;
|
||||
use super::norm::RMSNorm;
|
||||
use super::rope::RoPE;
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct DecoderLayer<B: Backend> {
|
||||
pub self_attn: Attention<B>,
|
||||
pub mlp: FeedForward<B>,
|
||||
pub input_layernorm: RmsNorm<B>,
|
||||
pub post_attention_layernorm: RmsNorm<B>,
|
||||
pub input_layernorm: RMSNorm<B>,
|
||||
pub post_attention_layernorm: RMSNorm<B>,
|
||||
}
|
||||
|
||||
impl<B: Backend> DecoderLayer<B> {
|
||||
@@ -27,8 +27,8 @@ impl<B: Backend> DecoderLayer<B> {
|
||||
) -> Self {
|
||||
let self_attn = Attention::new(hidden_size, n_heads, n_kv_heads, head_dim, device);
|
||||
let mlp = FeedForward::new(hidden_size, intermediate_size, device);
|
||||
let input_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let post_attention_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let input_layernorm = RMSNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let post_attention_layernorm = RMSNorm::new(hidden_size, rms_norm_eps, device);
|
||||
|
||||
Self {
|
||||
self_attn,
|
||||
|
||||
@@ -6,7 +6,7 @@ use burn::tensor::{Float, Int, Tensor};
|
||||
use super::attention::KVCache;
|
||||
use crate::config::LlamaConfig;
|
||||
use super::decoder::DecoderLayer;
|
||||
use super::norm::RmsNorm;
|
||||
use super::norm::RMSNorm;
|
||||
use super::rope::RoPE;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -26,7 +26,7 @@ impl<B: Backend> LlamaKVCache<B> {
|
||||
pub struct LlamaModel<B: Backend> {
|
||||
pub embed_tokens: Embedding<B>,
|
||||
pub layers: Vec<DecoderLayer<B>>,
|
||||
pub norm: RmsNorm<B>,
|
||||
pub norm: RMSNorm<B>,
|
||||
pub rope: RoPE,
|
||||
}
|
||||
|
||||
@@ -58,7 +58,7 @@ impl<B: Backend> LlamaForCausalLM<B> {
|
||||
));
|
||||
}
|
||||
|
||||
let norm = RmsNorm::new(config.hidden_size, config.rms_norm_eps, device);
|
||||
let norm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device);
|
||||
|
||||
let lm_head = LinearConfig::new(config.hidden_size, config.vocab_size).init(device);
|
||||
|
||||
|
||||
@@ -4,12 +4,12 @@ use burn::tensor::backend::Backend;
|
||||
use burn::tensor::{Float, Tensor};
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct RmsNorm<B: Backend> {
|
||||
pub struct RMSNorm<B: Backend> {
|
||||
pub weight: Param<Tensor<B, 1, Float>>,
|
||||
pub eps: f64,
|
||||
}
|
||||
|
||||
impl<B: Backend> RmsNorm<B> {
|
||||
impl<B: Backend> RMSNorm<B> {
|
||||
pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
|
||||
let weight = Initializer::Ones.init([dim], device);
|
||||
Self {
|
||||
|
||||
Reference in New Issue
Block a user