Machine Learning on Linera

The Linera application contract / service split allows for securely and efficiently running machine learning models on the edge.

The application's contract retrieves the correct model with all the correctness guarantees enforced by the consensus algorithm, while the client performs inference off-chain, in the un-metered service. Since the service is running on the user's own hardware, it can be implicitly trusted.

Guidelines

The existing examples use the candle framework by Hugging Face as the underlying ML framework.

candle is a minimalist ML framework for Rust with a focus on performance and usability. It also compiles to Wasm and has great support for Wasm both in and outside the browser. Check candle's examples for inspiration on the types of models which are supported.

Getting Started

To add ML capabilities to your existing Linera project, you'll need to add the candle-core, getrandom, rand and tokenizers dependencies to your Linera project:

candle-core = "0.4.1"
getrandom = { version = "0.2.12", default-features = false, features = ["custom"] }
rand = "0.8.5"

Optionally, to run Large Language Models, you'll also need the candle-transformers and transformers crate:

candle-transformers = "0.4.1"
tokenizers = { git = "https://github.com/christos-h/tokenizers", default-features = false, features = ["unstable_wasm"] }

Providing Randomness

ML frameworks use random numbers to perform inference. Linera services run in a Wasm VM which do not have access to the OS Rng. For this reason, we need to manually seed RNG used by candle. We do this by writing a custom getrandom.

Create a file under src/random.rs and add the following:

use std::sync::{Mutex, OnceLock};

use rand::{rngs::StdRng, Rng, SeedableRng};

static RNG: OnceLock<Mutex<StdRng>> = OnceLock::new();

fn custom_getrandom(buf: &mut [u8]) -> Result<(), getrandom::Error> {
    let seed = [0u8; 32];
    RNG.get_or_init(|| Mutex::new(StdRng::from_seed(seed)))
        .lock()
        .expect("failed to get RNG lock")
        .fill(buf);
    Ok(())
}

getrandom::register_custom_getrandom!(custom_getrandom);

This will enable candle and any other crates which rely on getrandom access to a deterministic RNG. If deterministic behaviour is not desired, the System API can be used to seed the RNG from a timestamp.

Loading the model into the Service

Models cannot currently be saved on-chain; for more information see the Limitations below.

To perform model inference, the model must be loaded into the service. To do this we'll use the fetch_url API when a query is made against the service:

impl Service for MyService {
    async fn handle_query(&self, request: Request) -> Response {
        // do some stuff here
        let raw_weights = self.runtime.fetch_url("https://my-model-provider.com/model.bin");
        // do more stuff here
    }
}

This can be served from a local webserver or pulled directly from a model provider such as Hugging Face.

At this point we have the raw bytes which correspond to the models and tokenizer. candle supports multiple formats for storing model weights, both quantized and not (gguf, ggml, safetensors, etc.).

Depending on the model format that you're using, candle exposes convenience functions to convert the bytes into a typed struct which can then be used to perform inference. Below is an example for a non-quantized Llama 2 model:

    fn load_llama_model(cursor: &mut Cursor<Vec<u8>>) -> Result<(Llama, Cache), candle_core::Error> {
        let config = llama2_c::Config::from_reader(cursor)?;
        let weights =
            llama2_c_weights::TransformerWeights::from_reader(cursor, &config, &Device::Cpu)?;
        let vb = weights.var_builder(&config, &Device::Cpu)?;
        let cache = llama2_c::Cache::new(true, &config, vb.pp("rot"))?;
        let llama = Llama::load(vb, config.clone())?;
        Ok((llama, cache))
    }

The addressable memory of the service's Wasm runtime being 4 GiB.
Not being able to load models directly to the GPU.

It is recommended that smaller models (50 Mb - 100 Mb) are used at current state of development.

Machine Learning on Linera

Guidelines

Getting Started

Providing Randomness

Loading the model into the Service

Inference

Limitations

Hardware Acceleration

On-Chain Models

Maximum Model Size