dnrops
/
candle
mirror of https://github.com/huggingface/candle.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make the metal sdpa tests deterministic. (#2750)

This commit is contained in:
Laurent Mazare 2025-01-28 09:05:24 +01:00 committed by GitHub
parent da02b59516
commit ab9019425a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
candle-nn/Cargo.toml
View File

@ -26,6 +26,7 @@ candle-metal-kernels = { workspace = true, optional = true }
anyhow = { workspace = true }
clap = { workspace = true }
rand = { workspace = true }
rand_distr = { workspace = true }
criterion = { workspace = true }
[features]
@ -37,4 +38,4 @@ metal = ["candle/metal", "dep:candle-metal-kernels", "dep:metal"]
[[bench]]
name = "bench_main"
harness = false
harness = false

123
candle-nn/tests/sdpa.rs
View File

@ -1,86 +1,84 @@
#[cfg(feature = "metal")]
mod metal_sdpa_tests {
#[test]
fn sdpa_full() -> candle::Result<()> {
use candle::{DType, Device, Tensor};
use candle::{DType, Device, Result, Shape, Tensor};
use rand::SeedableRng;
use rand_distr::Distribution;
use std::ops::{Div, Mul};
fn randn<S: Into<Shape>>(
rng: &mut rand::rngs::StdRng,
shape: S,
dev: &Device,
) -> Result<Tensor> {
let shape = shape.into();
let elem_count = shape.elem_count();
let normal = rand_distr::Normal::new(0.0, 1.0).unwrap();
let vs: Vec<f32> = (0..elem_count).map(|_| normal.sample(rng)).collect();
Tensor::from_vec(vs, &shape, dev)
}
#[test]
fn sdpa_full() -> Result<()> {
// Force seqlen = 100
const BS: usize = 4;
const R: usize = 4;
const L: usize = 4;
const DK: usize = 64;
const H: usize = 3;
let scale: f64 = f64::from(DK as u32).sqrt().recip();
let device = Device::new_metal(0)?;
let q = Tensor::randn(0f32, 1f32, (BS, H, R, DK), &device)?;
let k = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let v = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let mut rng = rand::rngs::StdRng::seed_from_u64(42);
let q = randn(&mut rng, (BS, H, R, DK), &device)?;
let k = randn(&mut rng, (BS, H, L, DK), &device)?;
let v = randn(&mut rng, (BS, H, L, DK), &device)?;
let ground_truth = {
let att = (q.clone() * scale)?.matmul(&k.clone().t()?)?;
let att = candle_nn::ops::softmax_last_dim(&att.to_dtype(DType::F32)?)?
.to_dtype(q.dtype())?;
att.matmul(&v.clone())?
};
let sdpa_output = candle_nn::ops::sdpa(&q, &k, &v, scale as f32, 1.)?;
assert_eq!(ground_truth.shape(), sdpa_output.shape());
let error: f32 = ((&ground_truth - &sdpa_output)?.abs()? / &ground_truth.abs()?)?
.sum_all()?
.to_scalar()?;
assert!(error <= 0.0005, "{}", error);
assert!(error <= 0.0004, "{}", error);
Ok(())
}
#[test]
fn sdpa_vector() -> candle::Result<()> {
use candle::{DType, Device, Tensor};
fn sdpa_vector() -> Result<()> {
// Allow vectorized, seqlen = 1
const BS: usize = 4;
const R: usize = 1;
const L: usize = 1;
const DK: usize = 64;
const H: usize = 3;
let scale: f64 = f64::from(DK as u32).sqrt().recip();
let device = Device::new_metal(0)?;
let q = Tensor::randn(0f32, 1f32, (BS, H, R, DK), &device)?;
let k = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let v = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let mut rng = rand::rngs::StdRng::seed_from_u64(4242);
let q = randn(&mut rng, (BS, H, R, DK), &device)?;
let k = randn(&mut rng, (BS, H, L, DK), &device)?;
let v = randn(&mut rng, (BS, H, L, DK), &device)?;
let ground_truth = {
let att = (q.clone() * scale)?.matmul(&k.clone().t()?)?;
let att = candle_nn::ops::softmax_last_dim(&att.to_dtype(DType::F32)?)?
.to_dtype(q.dtype())?;
att.matmul(&v.clone())?
};
let sdpa_output = candle_nn::ops::sdpa(&q, &k, &v, scale as f32, 1.)?;
assert_eq!(ground_truth.shape(), sdpa_output.shape());
let error: f32 = ((&ground_truth - &sdpa_output)?.abs()? / &ground_truth.abs()?)?
.sum_all()?
.to_scalar()?;
assert!(error <= 0.0001, "{}", error);
assert!(error <= 0.000, "{}", error);
Ok(())
}
#[test]
fn sdpa_full_softcapping() -> candle::Result<()> {
use candle::{DType, Device, Tensor};
use std::ops::{Div, Mul};
fn sdpa_full_softcapping() -> Result<()> {
// Allow vectorized, seqlen = 1
const BS: usize = 4;
const R: usize = 4;
@ -88,14 +86,13 @@ mod metal_sdpa_tests {
const DK: usize = 64;
const H: usize = 3;
const SOFTCAP: f64 = 50.;
let scale: f64 = f64::from(DK as u32).sqrt().recip();
let device = Device::new_metal(0)?;
let q = Tensor::randn(0f32, 1f32, (BS, H, R, DK), &device)?;
let k = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let v = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let mut rng = rand::rngs::StdRng::seed_from_u64(424242);
let q = randn(&mut rng, (BS, H, R, DK), &device)?;
let k = randn(&mut rng, (BS, H, L, DK), &device)?;
let v = randn(&mut rng, (BS, H, L, DK), &device)?;
let ground_truth = {
let att = (q.clone() * scale)?.matmul(&k.clone().t()?)?;
let att = candle_nn::ops::softmax_last_dim(
@ -107,25 +104,17 @@ mod metal_sdpa_tests {
.to_dtype(q.dtype())?;
att.matmul(&v.clone())?
};
let sdpa_output = candle_nn::ops::sdpa(&q, &k, &v, scale as f32, SOFTCAP as f32)?;
assert_eq!(ground_truth.shape(), sdpa_output.shape());
let error: f32 = ((&ground_truth - &sdpa_output)?.abs()? / &ground_truth.abs()?)?
.sum_all()?
.to_scalar()?;
assert!(error <= 0.0005, "{}", error);
Ok(())
}
#[test]
fn sdpa_vector_softcapping() -> candle::Result<()> {
use candle::{DType, Device, Tensor};
use std::ops::{Div, Mul};
fn sdpa_vector_softcapping() -> Result<()> {
// Allow vectorized, seqlen = 1
const BS: usize = 4;
const R: usize = 1;
@ -133,14 +122,13 @@ mod metal_sdpa_tests {
const DK: usize = 64;
const H: usize = 3;
const SOFTCAP: f64 = 50.;
let scale: f64 = f64::from(DK as u32).sqrt().recip();
let device = Device::new_metal(0)?;
let q = Tensor::randn(0f32, 1f32, (BS, H, R, DK), &device)?;
let k = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let v = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let mut rng = rand::rngs::StdRng::seed_from_u64(42424242);
let q = randn(&mut rng, (BS, H, R, DK), &device)?;
let k = randn(&mut rng, (BS, H, L, DK), &device)?;
let v = randn(&mut rng, (BS, H, L, DK), &device)?;
let ground_truth = {
let att = (q.clone() * scale)?.matmul(&k.clone().t()?)?;
let att = candle_nn::ops::softmax_last_dim(
@ -152,55 +140,42 @@ mod metal_sdpa_tests {
.to_dtype(q.dtype())?;
att.matmul(&v.clone())?
};
let sdpa_output = candle_nn::ops::sdpa(&q, &k, &v, scale as f32, SOFTCAP as f32)?;
assert_eq!(ground_truth.shape(), sdpa_output.shape());
let error: f32 = ((&ground_truth - &sdpa_output)?.abs()? / &ground_truth.abs()?)?
.sum_all()?
.to_scalar()?;
assert!(error <= 0.0001, "{}", error);
Ok(())
}
#[test]
fn sdpa_vector_cross() -> candle::Result<()> {
use candle::{DType, Device, Tensor};
fn sdpa_vector_cross() -> Result<()> {
// Allow vectorized, seqlen = 1. Simulat cross attention case where R != L, R = 1
const BS: usize = 4;
const R: usize = 1;
const L: usize = 24;
const DK: usize = 64;
const H: usize = 3;
let scale: f64 = f64::from(DK as u32).sqrt().recip();
let device = Device::new_metal(0)?;
let q = Tensor::randn(0f32, 1f32, (BS, H, R, DK), &device)?;
let k = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let v = Tensor::randn(0f32, 1f32, (BS, H, L, DK), &device)?;
let mut rng = rand::rngs::StdRng::seed_from_u64(4242424242);
let q = randn(&mut rng, (BS, H, R, DK), &device)?;
let k = randn(&mut rng, (BS, H, L, DK), &device)?;
let v = randn(&mut rng, (BS, H, L, DK), &device)?;
let ground_truth = {
let att = (q.clone() * scale)?.matmul(&k.clone().t()?)?;
let att = candle_nn::ops::softmax_last_dim(&att.to_dtype(DType::F32)?)?
.to_dtype(q.dtype())?;
att.matmul(&v.clone())?
};
let sdpa_output = candle_nn::ops::sdpa(&q, &k, &v, scale as f32, 1.)?;
assert_eq!(ground_truth.shape(), sdpa_output.shape());
let error: f32 = ((&ground_truth - &sdpa_output)?.abs()? / &ground_truth.abs()?)?
.sum_all()?
.to_scalar()?;
assert!(error <= 0.0013, "{}", error);
Ok(())
}
}