huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use a common with_tracing::RmsNorm in a few models. (#1871)

Browse Source 
* Add RmsNorm with tracing.

* Use with_tracing::RmsNorm in some models.
This commit is contained in:
Jani Monoses
2024-03-18 22:40:06 +02:00
committed by GitHub
parent 6a966cf9e0
commit 90fc82211f
6 changed files with 29 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
candle-transformers/src/models/yi.rs
View File

@ -1,5 +1,5 @@
/// https://huggingface.co/01-ai/Yi-6B/blob/main/modeling_yi.py
use crate::models::with_tracing::{linear_no_bias, Linear};
use crate::models::with_tracing::{linear_no_bias, Linear, RmsNorm};
use candle::{DType, Device, Module, Result, Tensor, D};
use candle_nn::{Activation, VarBuilder};
use std::sync::Arc;
@ -50,27 +50,6 @@ impl Config {
}
}
#[derive(Debug, Clone)]
struct RmsNorm {
inner: candle_nn::RmsNorm,
span: tracing::Span,
}
impl RmsNorm {
fn new(size: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
let span = tracing::span!(tracing::Level::TRACE, "rms-norm");
let inner = candle_nn::rms_norm(size, eps, vb)?;
Ok(Self { inner, span })
}
}
impl Module for RmsNorm {
fn forward(&self, x: &Tensor) -> Result<Tensor> {
let _enter = self.span.enter();
self.inner.forward(x)
}
}
#[derive(Debug, Clone)]
struct RotaryEmbedding {
sin: Tensor,