Tests migration to tch

CONTRIBUTING.md

@@ -131,6 +131,8 @@ If that looks like this, you are good to go 🎉
 
 ### Step 3: Run Tests
 
+#### PYO3 tests
+
 Finally, run the tests for the package using Cargo:
 
 ```sh
@@ -143,6 +145,26 @@ To run the `PyO3` tests, add the `pyo3` flag:
 cargo test --features pyo3
 ```
 
+#### `tch` tests
+
+`tch` based tests are ran behind the `testing` feature. You need to first have the PyTorch library (libtorch) in v2.3.0 to be available on your system. Follow the [official `tch` for more details](https://github.com/LaurentMazare/tch-rs/tree/main?tab=readme-ov-file). We'll use the libtorch library installed in the python envionment:
+
+```sh
+export LIBTORCH_USE_PYTORCH=1
+```
+
+You can now run tests:
+
+```sh
+cargo test --features testing
+```
+
+**NOTE**: If you're having compilation issue with MacOS. You can add the `libtorch` lib to your environment :
+
+```sh
+export DYLD_LIBRARY_PATH=$PWD/venv/lib/python3.10/site-packages/torch/lib:$DYLD_LIBRARY_PATH
+```
+
 ### Step 5: Run WASM Tests
 
 To run WASM tests (e.g., the whisper test) run:

Cargo.toml

@@ -5,8 +5,8 @@ members = [
     "crates/ratchet-web",
     "crates/ratchet-loader",
     "crates/ratchet-models",
-    "crates/ratchet-nn", 
-    "crates/ratchet-hub", 
+    "crates/ratchet-nn",
+    "crates/ratchet-hub",
     "crates/ratchet-cli",
 ]
 resolver = "2"
@@ -18,26 +18,30 @@ debug-assertions = true
 [profile.release]
 panic = 'abort'
 lto = "fat"
-codegen-units = 1 
+codegen-units = 1
 
 [profile.profiling]
 inherits = "release"
 debug = 2
 
 [workspace.dependencies]
 wgpu = { version = "0.20", features = ["fragile-send-sync-non-atomic-wasm"] }
-bytemuck = { version = "1.14.0", features=["wasm_simd", "aarch64_simd", "extern_crate_alloc"] }
+bytemuck = { version = "1.14.0", features = [
+    "wasm_simd",
+    "aarch64_simd",
+    "extern_crate_alloc",
+] }
 num-traits = "0.2.17"
 half = { version = "2.3.1", features = ["num-traits", "bytemuck"] }
 derive-new = "0.6.0"
 log = "0.4.20"
 thiserror = "1.0.56"
 byteorder = "1.5.0"
-npyz = { version = "0.8.3"}
+npyz = { version = "0.8.3" }
 hf-hub = "0.3.2"
 serde = "1.0"
 anyhow = "1.0.79"
-tokenizers = "0.19.1" 
+tokenizers = "0.19.1"
 
 js-sys = "0.3.64"
 wasm-bindgen = "0.2.91"
@@ -90,3 +94,4 @@ wasm-bindgen-futures = "0.4.41"
 web-sys = "0.3.64"
 web-time = "1.0.0"
 futures-intrusive = "0.5.0"
+tch = "0.16.0"

crates/ratchet-core/Cargo.toml

@@ -9,7 +9,7 @@ default = ["rand", "testing"]
 gpu-profiling = ["dep:tabled", "dep:itertools"]
 rand = ["dep:rand", "dep:rand_distr"]
 plotting = ["dep:dot3", "dep:tempfile"]
-testing = ["dep:npyz", "dep:ndarray"]
+testing = ["dep:npyz", "dep:ndarray","dep:tch"]
 pyo3 = ["dep:pyo3", "dep:numpy", "dep:regex"]
 
 [build-dependencies]
@@ -31,7 +31,7 @@ num-traits = { workspace = true }
 log = { workspace = true }
 thiserror = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
-anyhow.workspace = true 
+anyhow.workspace = true
 
 rustc-hash = { workspace = true }
 slotmap = { workspace = true }
@@ -55,12 +55,13 @@ tempfile = { workspace = true, optional = true }
 tabled = { workspace = true, optional = true }
 itertools = { workspace = true, optional = true }
 
-pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } 
+pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
 regex = { workspace = true, optional = true }
 numpy = { workspace = true, optional = true }
+tch = {workspace =true, optional=true}
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
-wasm-bindgen.workspace = true 
+wasm-bindgen.workspace = true
 futures-intrusive.workspace = true
 
 async-trait = "0.1.77"

crates/ratchet-core/src/ops/binary.rs

@@ -151,9 +151,9 @@ impl MetaOperation for Binary {
     }
 }
 
-#[cfg(all(test, feature = "pyo3"))]
+#[cfg(all(test, feature = "testing"))]
 mod tests {
-    use crate::{test_util::run_py_prg, BinaryOp, Device, DeviceRequest, Shape, Tensor};
+    use crate::{BinaryOp, Device, DeviceRequest, Shape, Tensor};
     use test_strategy::{proptest, Arbitrary};
 
     thread_local! {
@@ -168,16 +168,15 @@ mod tests {
     }
 
     fn ground_truth(a: &Tensor, b: &Tensor, op: &BinaryOp) -> anyhow::Result<Tensor> {
-        let kn = op.kernel_name();
-        let prg = format!(
-            r#"
-import torch
-def {}(a, b):
-    return torch.{}(torch.from_numpy(a), torch.from_numpy(b)).numpy()
-"#,
-            kn, kn
-        );
-        run_py_prg(prg.to_string(), &[a, b], &[])
+        let a = a.to_tch::<f32>()?;
+        let b = b.to_tch::<f32>()?;
+        let result = match op {
+            BinaryOp::Add => a.f_add(&b)?,
+            BinaryOp::Sub => a.f_sub(&b)?,
+            BinaryOp::Mul => a.f_mul(&b)?,
+            BinaryOp::Div => a.f_div(&b)?,
+        };
+        Tensor::try_from(&result)
     }
 
     fn run_binary_trial(prob: BinaryProblem) -> anyhow::Result<()> {

crates/ratchet-core/src/ops/matmul.rs

@@ -594,12 +594,10 @@ impl MetaOperation for GEMM {
     }
 }
 
-#[cfg(all(test, feature = "pyo3"))]
+#[cfg(all(test, feature = "testing"))]
 mod tests {
     use test_strategy::{proptest, Arbitrary};
 
-    use crate::test_util::run_py_prg;
-
     use crate::{shape, Device, DeviceRequest, Quantization, Quantizer};
 
     use super::*;
@@ -612,53 +610,22 @@ mod tests {
         trans_rhs: bool,
         trans_out: bool,
     ) -> anyhow::Result<Tensor> {
-        let a_op = if trans_lhs {
-            "torch.permute(torch.from_numpy(a), [0, 2, 1])"
-        } else {
-            "torch.from_numpy(a)"
-        };
-
-        let b_op = if trans_rhs {
-            "torch.permute(torch.from_numpy(b), [0, 2, 1])"
-        } else {
-            "torch.from_numpy(b)"
-        };
-
-        let inner = if bias.is_some() {
-            format!(
-                "torch.add(torch.matmul({}, {}), torch.from_numpy(bias))",
-                a_op, b_op
-            )
-        } else {
-            format!("torch.matmul({}, {})", a_op, b_op)
-        };
-
-        let result_op = if trans_out {
-            format!(
-                "np.ascontiguousarray(torch.permute({}, [0, 2, 1]).numpy())",
-                inner
-            )
-        } else {
-            format!("{}.numpy()", inner)
-        };
-
-        let prg = format!(
-            r#"
-import torch
-import numpy as np
-def matmul(a, b{}):
-    return {}"#,
-            if bias.is_some() { ", bias" } else { "" },
-            result_op
-        );
-
-        let args = if let Some(bias) = bias {
-            vec![a, b, bias]
-        } else {
-            vec![a, b]
+        let a = a.to_tch::<f32>()?;
+        let b = b.to_tch::<f32>()?;
+        let a = if trans_lhs { a.permute([0, 2, 1]) } else { a };
+        let b = if trans_rhs { b.permute([0, 2, 1]) } else { b };
+
+        let result = match bias {
+            Some(bias) => {
+                let bias = bias.to_tch::<f32>()?;
+                a.matmul(&b).f_add(&bias)?
+            }
+            None => a.matmul(&b),
         };
-
-        run_py_prg(prg.to_string(), &args, &[])
+        if trans_out {
+            return Tensor::try_from(&result.permute([0, 2, 1]).contiguous());
+        }
+        Tensor::try_from(&result)
     }
 
     #[derive(Arbitrary, Clone, Debug)]

crates/ratchet-core/src/ops/norm/groupnorm.rs

@@ -30,11 +30,10 @@ impl Operation for GroupNorm {
         Ok(self.norm.input.storage_view().clone())
     }
 }
-#[cfg(all(test, feature = "pyo3"))]
+#[cfg(all(test, feature = "testing"))]
 mod tests {
     use test_strategy::{proptest, Arbitrary};
 
-    use crate::test_util::run_py_prg;
     use crate::{rvec, shape, Device, DeviceRequest, Tensor};
 
     fn ground_truth(
@@ -43,20 +42,15 @@ mod tests {
         bias: Option<&Tensor>,
         num_groups: usize,
     ) -> anyhow::Result<Tensor> {
-        let prg = r#"
-import torch
-import torch.nn.functional as F
-
-def manual_group_norm(input, scale, bias, num_groups):
-    (input, scale, bias) = (torch.from_numpy(input), torch.from_numpy(scale), torch.from_numpy(bias))
-    return F.group_norm(input, num_groups, weight=scale, bias=bias).numpy()
-"#;
-
-        let inputs = match bias {
-            Some(bias) => rvec![input, scale, bias],
-            None => rvec![input, scale],
+        let input = input.to_tch::<f32>()?;
+        let scale = scale.to_tch::<f32>()?;
+        let bias = match bias {
+            Some(b) => Some(b.to_tch::<f32>()?),
+            None => None,
         };
-        run_py_prg(prg.to_string(), &inputs, &[&num_groups])
+        let result =
+            input.f_group_norm(num_groups as i64, Some(&scale), bias.as_ref(), 1e-5, false)?;
+        Tensor::try_from(&result)
     }
 
     fn run_norm_trial(device: &Device, problem: GroupNormProblem) -> anyhow::Result<()> {

crates/ratchet-core/src/ops/norm/mod.rs

@@ -180,11 +180,10 @@ impl MetaOperation for NormOp {
     }
 }
 
-#[cfg(all(test, feature = "pyo3"))]
+#[cfg(all(test, feature = "testing"))]
 mod tests {
     use test_strategy::{proptest, Arbitrary};
 
-    use crate::test_util::run_py_prg;
     use crate::{rvec, shape, Device, DeviceRequest, Tensor};
 
     fn ground_truth(
@@ -193,35 +192,33 @@ mod tests {
         scale: &Tensor,
         bias: Option<&Tensor>,
     ) -> anyhow::Result<Tensor> {
-        let ln_prg = r#"
-import torch
-import torch.nn.functional as F
-
-def layer_norm(input, scale, bias):
-    (input, scale, bias) = (torch.from_numpy(input), torch.from_numpy(scale), torch.from_numpy(bias))
-    return F.layer_norm(input, (input.shape[-1],), weight=scale, bias=bias).numpy()
-"#;
-
-        let rms_prg = r#"
-import torch
-def manual_rms_norm(input, scale):
-    (input, scale) = (torch.from_numpy(input), torch.from_numpy(scale))
-    variance = input.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
-    input = input * torch.rsqrt(variance + 1e-5)
-    return (scale * input).numpy()
-"#;
-
-        let prg = match var {
-            NormVariant::LayerNorm => ln_prg,
-            NormVariant::RMSNorm => rms_prg,
+        let input = input.to_tch::<f32>()?;
+        let scale = scale.to_tch::<f32>()?;
+        let bias = match bias {
+            Some(b) => Some(b.to_tch::<f32>()?),
+            None => None,
         };
-
-        let inputs = match bias {
-            Some(bias) => rvec![input, scale, bias],
-            None => rvec![input, scale],
+        let result = match var {
+            NormVariant::LayerNorm => input.f_layer_norm(
+                [*input.size().last().unwrap()],
+                Some(&scale),
+                bias.as_ref(),
+                1e-5,
+                false,
+            )?,
+            NormVariant::RMSNorm => {
+                // (input, scale) = (torch.from_numpy(input), torch.from_numpy(scale))
+                // variance = input.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+                // input = input * torch.rsqrt(variance + 1e-5)
+                // return (scale * input).numpy()
+                let variance = input
+                    .f_pow_tensor_scalar(2)?
+                    .mean_dim(-1, true, input.kind());
+                let input = input.multiply(&variance.f_add_scalar(1e-5)?.rsqrt());
+                scale.multiply(&input)
+            }
         };
-
-        run_py_prg(prg.to_string(), &inputs, &[])
+        Tensor::try_from(&result)
     }
 
     fn run_norm_trial(device: &Device, problem: NormProblem) -> anyhow::Result<()> {

crates/ratchet-core/src/ops/softmax.rs

@@ -104,25 +104,19 @@ impl MetaOperation for Softmax {
     }
 }
 
-#[cfg(all(test, feature = "pyo3"))]
+#[cfg(all(test, feature = "testing"))]
 mod tests {
-    use test_strategy::{proptest, Arbitrary};
-
-    use crate::test_util::run_py_prg;
     use crate::{shape, Device, DeviceRequest, Tensor};
+    use tch;
+    use test_strategy::{proptest, Arbitrary};
 
     thread_local! {
         static GPU_DEVICE: Device = Device::request_device(DeviceRequest::GPU).unwrap();
     }
 
     fn ground_truth(a: &Tensor) -> anyhow::Result<Tensor> {
-        let prg = r#"
-import torch
-import torch.nn.functional as F
-def softmax(a):
-    return F.softmax(torch.from_numpy(a), dim=-1).numpy()
-"#;
-        run_py_prg(prg.to_string(), &[a], &[])
+        let t = a.to_tch::<f32>()?;
+        Tensor::try_from(&t.softmax(-1, Some(tch::kind::Kind::Float)))
     }
 
     fn run_softmax_trial(problem: SoftmaxProblem) {