use crate::{Tensor, AutogradMetaT, Dim, FloatDType, Layout, Storage, StorageIndices, D};
use super::reduce::DimArray;
impl<T: FloatDType> Tensor<T> {
pub fn rms_norm(input: &Tensor<T>, weight: &Tensor<T>, eps: T) -> crate::Result<Tensor<T>> {
let variance = input.sqr()?.mean_keepdim(D::Minus1)?;
let rms = (variance + eps).sqrt()?;
let input_normalized = input.broadcast_div(&rms)?;
let out = input_normalized.broadcast_mul(weight)?;
Ok(out)
}
pub fn rms_norm_fused(input: &Tensor<T>, weight: &Tensor<T>, eps: T) -> crate::Result<Tensor<T>> {
let last_dim = D::Minus1.to_index(input.shape(), "rms_norm")?;
let input_layout = input.layout();
let last_dim_stride = input_layout.stride()[last_dim];
let last_dim_size = input_layout.dims()[last_dim];
let weight_size = weight.dims1()?;
if last_dim_size != weight_size {
crate::bail!("input dim size != weight size in rms norm");
}
let input_storage = input.storage_read()?;
let input_data = input_storage.data();
let weight_storage = weight.storage_read()?;
let weight_data = weight_storage.data();
let weight_data = &weight_data[weight.layout().start_offset..];
let weight_stride = weight.layout().stride()[0];
let mut output_data = vec![T::ZERO; input.element_count()];
let batch_size = output_data.len() / last_dim_size;
match input_layout.storage_indices() {
StorageIndices::Contiguous(index) => {
let input_data = &input_data[index.begin_index..index.end_index];
for b in 0..batch_size {
let input_data = &input_data[b*last_dim_size..b*last_dim_size+last_dim_size];
let output_data = &mut output_data[b*last_dim_size..b*last_dim_size+last_dim_size];
let variance = input_data.iter().cloned().map(|v| v.sqr()).sum::<T>() / T::from_usize(last_dim_size);
let rms = (variance + eps).sqrt();
for i in 0..last_dim_size {
output_data[i] = input_data[i] / rms * weight_data[i * weight_stride];
}
}
}
StorageIndices::Uncontiguous(_) => {
let layout = input_layout.narrow(last_dim, 0, 1)?;
println!("{}", layout.storage_indices().count());
for (b, src_index) in layout.storage_indices().enumerate() {
let output_data = &mut output_data[b*last_dim_size..b*last_dim_size+last_dim_size];
let input_arr: DimArray<'_, T> = DimArray {
src: &input_data[src_index..],
size: last_dim_size,
stride: last_dim_stride,
};
let variance = input_arr.clone().into_iter().map(|v| v.sqr()).sum::<T>() / T::from_usize(last_dim_size);
let rms = (variance + eps).sqrt();
for (i, v) in (0..last_dim_size).into_iter().zip(input_arr.into_iter()) {
output_data[i] = v / rms * weight_data[i * weight_stride];
}
}
}
}
let output_storage = Storage::new(output_data);
let meta = T::AutogradMeta::on_rms_norm_op(input, weight, eps);
let output = Tensor::from_storage(output_storage, input.shape(), meta);
Ok(output)
}
pub(crate) fn rms_norm_backward_fused(
grad_output: &Tensor<T>,
input: &Tensor<T>,
weight: &Tensor<T>,
input_sum_grad: &Tensor<T>,
weight_sum_grad: &Tensor<T>,
eps: T,
) -> crate::Result<()> {
let shape = input.shape();
let last_dim = shape.dims().len() - 1;
let last_dim_size = shape.dims()[last_dim];
let batch_size = input.element_count() / last_dim_size;
let d_inv = T::from_f64(1.0 / last_dim_size as f64);
let g_data_storage = grad_output.storage_read()?;
let g_data = g_data_storage.data();
assert!(grad_output.is_contiguous());
let x_data_storage = input.storage_read()?;
let x_data = x_data_storage.data();
let w_data_storage = weight.storage_read()?;
let w_data = w_data_storage.data();
let w_data = &w_data[weight.layout().start_offset..];
let w_stride = weight.layout().stride()[0];
assert!(input_sum_grad.is_contiguous());
assert!(weight_sum_grad.is_contiguous());
let mut input_grad_storage = input_sum_grad.storage_write()?;
let mut weight_gard_storage = weight_sum_grad.storage_write()?;
let input_gard_data = input_grad_storage.data_mut();
let weight_gard_data = weight_gard_storage.data_mut();
if input.is_contiguous() {
let x_data = &x_data[input.layout().start_offset..];
for b in 0..batch_size {
let offset = b * last_dim_size;
let g_row = &g_data[offset..offset + last_dim_size];
let x_row = &x_data[offset..offset + last_dim_size];
let input_gard_row_data = &mut input_gard_data[offset..offset + last_dim_size];
let mut sum_sq = T::ZERO;
let mut sum_dot = T::ZERO;
for i in 0..last_dim_size {
let xi = x_row[i];
let wi = w_data[i * w_stride];
let gi = g_row[i];
sum_sq = sum_sq + xi * xi;
sum_dot = sum_dot + gi * wi * xi;
}
let variance = sum_sq * d_inv + eps;
let rsqrt_v = T::one() / variance.sqrt(); let s_v = (sum_dot * d_inv) / variance;
for i in 0..last_dim_size {
let xi = x_row[i];
let wi = w_data[i * w_stride];
let gi = g_row[i];
input_gard_row_data[i] += rsqrt_v * (wi * gi - xi * s_v);
weight_gard_data[i] += gi * (xi * rsqrt_v);
}
}
} else {
unimplemented!("rms norm backward in un continue")
}
Ok(())
}
pub(crate) fn softmax_backward_fused(
grad_output: &Tensor<T>,
output: &Tensor<T>,
input_grad: &Tensor<T>,
dim: usize,
) -> crate::Result<()> {
let reduce_dim = dim;
let shape = output.shape();
let reduce_dim_size = shape.dims()[reduce_dim];
let element_count = output.layout().element_count();
let g_data_storage = grad_output.storage_read()?;
let g_data = g_data_storage.data();
let y_data_storage = output.storage_read()?;
let y_data = y_data_storage.data();
assert!(input_grad.is_contiguous());
let mut dx_storage = input_grad.storage_write()?;
let dx_data = dx_storage.data_mut();
let y_stride = output.layout().stride()[reduce_dim];
let g_stride = grad_output.layout().stride()[reduce_dim];
if output.is_contiguous() && grad_output.is_contiguous() && y_stride == 1 && g_stride == 1 {
let y_data = &y_data[output.layout().start_offset..output.layout().start_offset + element_count];
let g_data = &g_data[grad_output.layout().start_offset..grad_output.layout().start_offset + element_count];
let dx_data = &mut dx_data[input_grad.layout().start_offset..input_grad.layout().start_offset + element_count];
let batch_size = element_count / reduce_dim_size;
for b in 0..batch_size {
let offset = b * reduce_dim_size;
let g_row = &g_data[offset..offset + reduce_dim_size];
let y_row = &y_data[offset..offset + reduce_dim_size];
let dx_row = &mut dx_data[offset..offset + reduce_dim_size];
let mut sum_gy = T::ZERO;
for i in 0..reduce_dim_size {
sum_gy = sum_gy + g_row[i] * y_row[i];
}
for i in 0..reduce_dim_size {
let dy_i = g_row[i];
let y_i = y_row[i];
dx_row[i] += y_i * (dy_i - sum_gy);
}
}
} else {
unimplemented!("softmax backward in non-contiguous memory is not implemented yet")
}
Ok(())
}
}
impl<T: FloatDType> Tensor<T> {
pub fn softmax<D: Dim>(&self, dim: D) -> crate::Result<Tensor<T>> {
let dim = dim.to_index(self.shape(), "softmax")?;
let max = self.max_keepdim(dim)?; let diff = self.broadcast_sub(&max)?; let num = diff.exp()?; let den = num.sum_keepdim(dim)?;
let out = num.broadcast_div(&den)?;
Ok(out)
}
pub fn softmax_fused<D: Dim>(&self, dim: D) -> crate::Result<Tensor<T>> {
let reduce_dim = dim.to_index(self.shape(), "softmax_fused")?;
assert!(reduce_dim < self.layout().dims().len());
let reduce_dim_stride = self.layout().stride()[reduce_dim];
let reduce_dim_size = self.layout().dims()[reduce_dim];
let src_storage = self.storage_read()?;
let src_data = src_storage.data();
let src_start_offset = self.layout().start_offset;
let element_count = self.layout().element_count();
let mut dst: Vec<T> = Vec::with_capacity(element_count);
unsafe { dst.set_len(element_count) };
if self.is_contiguous() && reduce_dim_stride == 1 {
let src_data = &src_data[src_start_offset..src_start_offset + element_count];
for (src_chunk, dst_chunk) in
src_data.chunks(reduce_dim_size).zip(dst.chunks_mut(reduce_dim_size))
{
let mut m = src_chunk[0];
let mut s = T::one();
for &x in src_chunk.iter().skip(1) {
if x <= m {
s = s + (x - m).exp();
} else {
s = s * (m - x).exp() + T::one();
m = x;
}
}
for (i, &x) in src_chunk.iter().enumerate() {
dst_chunk[i] = (x - m).exp() / s;
}
}
} else {
let input_layout = self.layout().narrow(reduce_dim, 0, 1)?;
let out_layout = Layout::contiguous(self.shape());
let out_layout_narrow = out_layout.narrow(reduce_dim, 0, 1)?;
let out_stride = out_layout.stride()[reduce_dim];
for (src_offset, dst_offset) in
input_layout.storage_indices().zip(out_layout_narrow.storage_indices())
{
let mut m = src_data[src_offset];
let mut s = T::one();
for i in 1..reduce_dim_size {
let x = src_data[src_offset + i * reduce_dim_stride];
if x <= m {
s = s + (x - m).exp();
} else {
s = s * (m - x).exp() + T::one();
m = x;
}
}
for i in 0..reduce_dim_size {
let x = src_data[src_offset + i * reduce_dim_stride];
dst[dst_offset + i * out_stride] = (x - m).exp() / s;
}
}
}
let storage = Storage::new(dst);
let meta = T::AutogradMeta::on_softmax_op(self, reduce_dim);
Ok(Tensor::<T>::from_storage(storage, self.shape(), meta))
}
}
#[cfg(test)]
mod tests {
use crate::{Tensor, Var};
#[test]
fn test_rms_norm_forward() -> crate::Result<()> {
let weight = Tensor::randn(0.0, 1.0, (64,))?;
const EPS: f64 = 1e-4;
let input = Tensor::randn(0.0, 1.0, (4, 10, 64))?;
println!("{}", input.is_contiguous());
let result1 = Tensor::rms_norm(&input, &weight, EPS)?;
let result_fused1 = Tensor::rms_norm_fused(&input, &weight, EPS)?;
assert!(result1.allclose(&result_fused1, 1e-5, 8e-8)?);
let input = input.transpose_last()?.copy()?.transpose_last()?;
println!("{}", input.is_contiguous());
let result2 = Tensor::rms_norm(&input, &weight, EPS)?;
let result_fused2 = Tensor::rms_norm_fused(&input, &weight, EPS)?;
assert!(result1.allclose(&result2, 1e-5, 8e-8)?);
assert!(result1.allclose(&result_fused2, 1e-5, 8e-8)?);
Ok(())
}
use std::time::Instant;
#[test]
fn test_rms_norm_backward() -> crate::Result<()> {
let input = Var::randn(0.0, 1.0, (1, 32, 128))?;
let weight = Var::randn(0.0, 1.0, (128,))?;
const EPS: f64 = 1e-4;
{
let _ = Tensor::rms_norm(&input, &weight, EPS)?;
let _ = Tensor::rms_norm_fused(&input, &weight, EPS)?;
}
let start = Instant::now();
let result = Tensor::rms_norm(&input, &weight, EPS)?;
let forward_time_1 = start.elapsed();
let start = Instant::now();
let grads1 = result.backward()?;
let backward_time_1 = start.elapsed();
let start = Instant::now();
let result_fused = Tensor::rms_norm_fused(&input, &weight, EPS)?;
let forward_time_2 = start.elapsed();
let start = Instant::now();
let grads2 = result_fused.backward()?;
let backward_time_2 = start.elapsed();
assert!(result.allclose(&result_fused, 1e-5, 8e-8)?);
assert!(grads1[&input].allclose(&grads2[&input], 1e-5, 8e-8)?);
assert!(grads1[&weight].allclose(&grads2[&weight], 1e-5, 8e-8)?);
println!("=== Forward ===");
println!("rms_norm : {:?}", forward_time_1);
println!("rms_norm_fused : {:?}", forward_time_2);
println!("=== Backward ===");
println!("rms_norm : {:?}", backward_time_1);
println!("rms_norm_fused : {:?}", backward_time_2);
println!("=== Speedup ===");
println!(
"forward speedup : {:.2}x",
forward_time_1.as_secs_f64() / forward_time_2.as_secs_f64()
);
println!(
"backward speedup: {:.2}x",
backward_time_1.as_secs_f64() / backward_time_2.as_secs_f64()
);
Ok(())
}
#[test]
fn test_softmax_forward() -> crate::Result<()> {
let input = Tensor::randn(0.0, 1.0, (4, 10, 64))?;
println!("{}", input.is_contiguous());
let result1 = Tensor::softmax(&input, 2)?;
let result_fused1 = Tensor::softmax_fused(&input, 2)?;
assert!(result1.allclose(&result_fused1, 1e-5, 8e-8)?);
let input = input.transpose_last()?.copy()?.transpose_last()?;
println!("{}", input.is_contiguous());
let result2 = Tensor::softmax(&input, 2)?;
let result_fused2 = Tensor::softmax_fused(&input, 2)?;
assert!(result1.allclose(&result2, 1e-5, 8e-8)?);
assert!(result1.allclose(&result_fused2, 1e-5, 8e-8)?);
Ok(())
}
#[test]
fn test_softmax_forward_not_last_dim() -> crate::Result<()> {
let input = Tensor::randn(0.0, 1.0, (4, 10, 64))?;
println!("{}", input.is_contiguous());
let result1 = Tensor::softmax(&input, 1)?;
let result_fused1 = Tensor::softmax_fused(&input, 1)?;
assert!(result1.allclose(&result_fused1, 1e-5, 8e-8)?);
let input = input.transpose_last()?.copy()?.transpose_last()?;
println!("{}", input.is_contiguous());
let result2 = Tensor::softmax(&input, 1)?;
let result_fused2 = Tensor::softmax_fused(&input, 1)?;
assert!(result1.allclose(&result2, 1e-5, 8e-8)?);
assert!(result1.allclose(&result_fused2, 1e-5, 8e-8)?);
Ok(())
}
#[test]
fn test_softmax_backward() -> crate::Result<()> {
let shape = (8, 32, 128);
let dims_to_test = [2];
for dim in dims_to_test {
println!("\n========================================");
println!("Testing Softmax along dim: {}", dim);
println!("========================================");
let input = Var::randn(0.0, 1.0, shape)?;
{
let _ = input.softmax(dim)?;
let _ = input.softmax_fused(dim)?;
}
let start = Instant::now();
let result = input.softmax(dim)?;
let forward_time_1 = start.elapsed();
let start = Instant::now();
let grads1 = result.backward()?;
let backward_time_1 = start.elapsed();
let start = Instant::now();
let result_fused = input.softmax_fused(dim)?;
let forward_time_2 = start.elapsed();
let start = Instant::now();
let grads2 = result_fused.backward()?;
let backward_time_2 = start.elapsed();
assert!(
result.allclose(&result_fused, 1e-5, 8e-8)?,
"Forward mismatch at dim {}", dim
);
assert!(
grads1[&input].allclose(&grads2[&input], 1e-5, 8e-8)?,
"Backward mismatch at dim {}", dim
);
println!("=== Forward ===");
println!("softmax : {:?}", forward_time_1);
println!("softmax_fused : {:?}", forward_time_2);
println!("=== Backward ===");
println!("softmax : {:?}", backward_time_1);
println!("softmax_fused : {:?}", backward_time_2);
println!("=== Speedup ===");
println!(
"forward speedup : {:.2}x",
forward_time_1.as_secs_f64() / forward_time_2.as_secs_f64()
);
println!(
"backward speedup: {:.2}x",
backward_time_1.as_secs_f64() / backward_time_2.as_secs_f64()
);
}
Ok(())
}
}