mirror of
https://github.com/huggingface/candle.git
synced 2025-06-17 02:58:50 +00:00
af7f8b87d3
* Trying out a custom RmsNorm cuda kernel. * CPU implementation for rms-norm. * Cuda wrappers. * Add some validation. * Add some testing. * More testing.
91 lines
2.9 KiB
Rust
91 lines
2.9 KiB
Rust
#[cfg(feature = "mkl")]
|
|
extern crate intel_mkl_src;
|
|
|
|
#[cfg(feature = "accelerate")]
|
|
extern crate accelerate_src;
|
|
|
|
use candle::{test_device, test_utils::to_vec3_round, Device, Result, Tensor};
|
|
|
|
fn softmax(device: &Device) -> Result<()> {
|
|
let data = &[[[3f32, 1., 4.], [1., 5., 9.]], [[2., 1., 7.], [8., 2., 8.]]];
|
|
let tensor = Tensor::new(data, device)?;
|
|
let t0 = candle_nn::ops::softmax(&tensor.log()?, 0)?;
|
|
let t1 = candle_nn::ops::softmax(&tensor.log()?, 1)?;
|
|
let t2 = candle_nn::ops::softmax(&tensor.log()?, 2)?;
|
|
assert_eq!(
|
|
to_vec3_round(&t0, 4)?,
|
|
&[
|
|
// 3/5, 1/2, 4/11
|
|
[[0.6, 0.5, 0.3636], [0.1111, 0.7143, 0.5294]],
|
|
// 2/5, 1/2, 7/11
|
|
[[0.4, 0.5, 0.6364], [0.8889, 0.2857, 0.4706]]
|
|
]
|
|
);
|
|
assert_eq!(
|
|
to_vec3_round(&t1, 4)?,
|
|
&[
|
|
// 3/4, 1/6, 4/13
|
|
[[0.75, 0.1667, 0.3077], [0.25, 0.8333, 0.6923]],
|
|
// 2/10, 1/3, 7/15
|
|
[[0.2, 0.3333, 0.4667], [0.8, 0.6667, 0.5333]]
|
|
]
|
|
);
|
|
assert_eq!(
|
|
to_vec3_round(&t2, 4)?,
|
|
&[
|
|
// (3, 1, 4) / 8, (1, 5, 9) / 15
|
|
[[0.375, 0.125, 0.5], [0.0667, 0.3333, 0.6]],
|
|
// (2, 1, 7) / 10, (8, 2, 8) / 18
|
|
[[0.2, 0.1, 0.7], [0.4444, 0.1111, 0.4444]]
|
|
]
|
|
);
|
|
let t2 = candle_nn::ops::softmax_last_dim(&tensor.log()?)?;
|
|
assert_eq!(
|
|
to_vec3_round(&t2, 4)?,
|
|
&[
|
|
// (3, 1, 4) / 8, (1, 5, 9) / 15
|
|
[[0.375, 0.125, 0.5], [0.0667, 0.3333, 0.6]],
|
|
// (2, 1, 7) / 10, (8, 2, 8) / 18
|
|
[[0.2, 0.1, 0.7], [0.4444, 0.1111, 0.4444]]
|
|
]
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
fn rms_norm(device: &Device) -> Result<()> {
|
|
let data = &[[[3f32, 1., 4.], [1., 5., 9.]], [[2., 1., 7.], [8., 2., 8.]]];
|
|
let tensor = Tensor::new(data, device)?;
|
|
let alpha = Tensor::new(&[1f32, 2f32, 3f32], device)?;
|
|
let t = candle_nn::ops::rms_norm(&tensor, &alpha, 1e-5)?;
|
|
assert_eq!(
|
|
to_vec3_round(&t, 4)?,
|
|
&[
|
|
[[1.019, 0.6794, 4.0762], [0.1674, 1.6744, 4.521]],
|
|
[[0.4714, 0.4714, 4.9497], [1.206, 0.603, 3.6181]]
|
|
]
|
|
);
|
|
let t2 = candle_nn::ops::rms_norm_slow(&tensor, &alpha, 1e-5)?;
|
|
assert_eq!(
|
|
to_vec3_round(&t2, 4)?,
|
|
&[
|
|
[[1.019, 0.6794, 4.0762], [0.1674, 1.6744, 4.521]],
|
|
[[0.4714, 0.4714, 4.9497], [1.206, 0.603, 3.6181]]
|
|
]
|
|
);
|
|
let diff = (t - t2)?.abs()?.sum_all()?.to_vec0::<f32>()?;
|
|
assert!(diff < 1e-5);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn softmax_numerical_stability() -> Result<()> {
|
|
let dev = &Device::Cpu;
|
|
let xs = Tensor::new(&[1234f32, 0.], dev)?;
|
|
let softmax = candle_nn::ops::softmax(&xs, 0)?;
|
|
assert_eq!(softmax.to_vec1::<f32>()?, &[1f32, 0.]);
|
|
Ok(())
|
|
}
|
|
|
|
test_device!(softmax, softmax_cpu, softmax_gpu, softmax_metal);
|
|
test_device!(rms_norm, rms_norm_cpu, rms_norm_gpu, rms_norm_metal);
|