Support GGML q4_0 and q4_1 quantization.

README.md

@@ -1,5 +1,13 @@
 # Smol Rust RWKV
 
+## Dev Branch Note
+
+This branch currently needs a patched version of `ggml` and `ggml-sys` from the `llama-rs` project.
+
+You can find it here: https://github.com/KerfuffleV2/llama-rs/tree/experiment-ggml-map-ops
+
+`smolrwkv/Cargo.toml` is set up to look for it checked out in `../llama-rs` relative to this repo's directory.
+
 ## What is it?
 
 A simple example of the RWKV approach to language models written in Rust by someone that

smolrwkv-cli/src/args.rs

@@ -10,11 +10,23 @@ const DEFAULT_MODEL: &str = "./RWKV-4-Pile-430M-20220808-8066.safetensors";
 const DEFAULT_TOKENIZER: &str = "./20B_tokenizer.json";
 
 #[derive(Debug, Clone, Copy, ValueEnum)]
-#[value(rename_all = "lower")]
 pub enum EvalType {
+    #[value(name = "ndf32")]
+    /// ndarray-backed 32 bit floats. Uses a lot of memory.
     NDf32,
+    #[value(name = "ndq8")]
+    /// ndarray-backed 8 bit quantized. Better memory usage but quite slow.
     NDu8,
+    #[value(name = "ggmlf32")]
+    /// GGML-backed 32 bit. As above, uses a lot of memory.
     GGMLf32,
+    #[value(name = "ggmlq4_0")]
+    /// GGML-backed 4 bit quantized, method 1. Poor quality.
+    GGMLQ4_0,
+    #[value(name = "ggmlq4_1")]
+    /// GGML-backed 4 bit quantized, method 2. Decenent quality,
+    /// but slower (to load?)
+    GGMLQ4_1,
 }
 
 #[derive(Clone, Debug, Parser)]

smolrwkv-cli/src/main.rs

@@ -59,7 +59,7 @@ fn go() -> Result<()> {
         args::EvalType::NDf32 => {
             Ctx::NdFloat32(run_threadlimited(args.max_load_threads, move || {
                 anyhow::Ok({
-                    info!("Model type: non-quantized (full 32bit).");
+                    info!("Backend type: NDArray non-quantized (full 32bit).");
                     S::context::RWKVContext::<FloatType, Array2<FloatType>>::new(
                         tdm.try_into()?,
                         tokenizer,
@@ -71,16 +71,22 @@ fn go() -> Result<()> {
         args::EvalType::NDu8 => {
             Ctx::NdQuant8(run_threadlimited(args.max_load_threads, move || {
                 anyhow::Ok({
-                    info!("Model type: 8 bit-quantized weights.");
+                    info!("Backend type: NDArray 8 bit-quantized weights.");
                     S::context::RWKVContext::<FloatType, TensorQ2>::new(tdm.try_into()?, tokenizer)
                 })
             })?)
         }
-        args::EvalType::GGMLf32 => {
-            use smolrwkv::ggml::context::RWKVContext;
-            info!("Model type: GGML 32 bit.");
+        args::EvalType::GGMLf32 | args::EvalType::GGMLQ4_0 | args::EvalType::GGMLQ4_1 => {
+            use smolrwkv::ggml::{context::RWKVContext, loader::RwkvGgmlType};
+            let wtype = match args.eval_mode {
+                args::EvalType::GGMLf32 => RwkvGgmlType::Float32,
+                args::EvalType::GGMLQ4_0 => RwkvGgmlType::Q4_0,
+                args::EvalType::GGMLQ4_1 => RwkvGgmlType::Q4_1,
+                _ => panic!("Impossible: Bad eval mode!"),
+            };
+            info!("Backend type: GGML {wtype:?}");
             Ctx::GgmlFloat32(RWKVContext::new(
-                tdm.try_into()?,
+                (wtype, tdm).try_into()?,
                 tokenizer,
                 args.max_eval_threads,
             ))

smolrwkv-cli/src/util.rs

@@ -1,3 +1,4 @@
+#![allow(dead_code)]
 use std::io::Write;
 
 use anyhow::Result;

smolrwkv/Cargo.toml

@@ -31,3 +31,6 @@ optional = true
 
 [dependencies.ggml]
 path = "../../llama-rs/ggml"
+
+[dependencies.ggml-sys]
+path = "../../llama-rs/ggml-sys"