python sentence transformer all-MiniLM-L6-v2 is almost 2x faster than candle

[AbhishekBose]

Here is my candle implementation: (Taken from the examples itself)

`pub fn encode(&self, prompt: &str) -> Result<(Tensor,Tensor)> {

    let tokens = self.tokenizer
        .encode(prompt, true)
        .map_err(E::msg)?
        .get_ids()
        .to_vec();
    let token_ids = Tensor::new(&tokens[..],&self.device )?.unsqueeze(0)?;
    let token_type_ids = token_ids.zeros_like()?;
    let embeddings =self.model.forward(&token_ids, &token_type_ids)?;
    Ok((embeddings,token_ids))
}`

and here is the python implementation

model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2') embeddings = model.encode([request.query])

On comparing just the encoding times, this is what I get

Encoding time for rust:

Time taken for encoding: 57.054333ms
Time taken for encoding: 59.913916ms
Time taken for encoding: 55.118625ms
Time taken for encoding: 51.580917ms
Time taken for encoding: 60.823625ms
Time taken for encoding: 56.318333ms
Time taken for encoding: 52.357875ms
Time taken for encoding: 82.0645ms
Time taken for encoding: 52.349709ms
Time taken for encoding: 63.768209ms
Time taken for encoding: 55.508666ms

Encoding time for python:

Time taken for encoding: 33.95 ms
Time taken for encoding: 124.68 ms
Time taken for encoding: 54.5 ms
Time taken for encoding: 30.46 ms
Time taken for encoding: 20.73 ms
Time taken for encoding: 26.07 ms
Time taken for encoding: 37.49 ms
Time taken for encoding: 24.42 ms
Time taken for encoding: 36.08 ms
Time taken for encoding: 24.55 ms
Time taken for encoding: 36.13 ms
Time taken for encoding: 29.97 ms
Time taken for encoding: 35.69 ms
Time taken for encoding: 26.8 ms
Time taken for encoding: 31.32 ms
Time taken for encoding: 30.12 ms
Time taken for encoding: 32.37 ms
Time taken for encoding: 34.27 ms
Time taken for encoding: 31.85 ms
Time taken for encoding: 35.78 ms
Time taken for encoding: 44.09 ms
Time taken for encoding: 19.15 ms
Time taken for encoding: 23.83 ms
Time taken for encoding: 33.09 ms
Time taken for encoding: 31.65 ms

Is it the issue of the implementation or am I doing something wrong here?

The experiment was run on an m1 mac air

[BerserkerMother]

I have played with this, running Rust's release binary achieves faster results than Python's code; however, Python is about 2x faster than debug binary. Maybe there is something happening with JIT. Could you try your program in release binary?

[AbhishekBose]

These numbers are for the release binary. Wondering if there's something off with the rust implementation

[BerserkerMother]

I think you should submit the issue on tokenizers too. Python implementation is using rust bindings, no?
Could U provide the full example to replicate the problem? Have you used encode_batch?

[AbhishekBose]

@BerserkerMother
I think I am gonna do that. I was under the impression that sentence-transformers does use rust bindings.

This is the language model struct which loads and contains the encode method

impl LanguageModel {
    pub fn build_model_and_tokenizer() -> Result<LanguageModel> {
       let device = device(true)?;
        let default_model = "sentence-transformers/all-MiniLM-L6-v2".to_string();
        let default_revision = "refs/pr/21".to_string();
        let default_model_id = None;
        let revision = None;
        let (model_id, revision) = match (default_model_id, revision) {
            (Some(model_id), Some(revision)) => (model_id, revision),
            (Some(model_id), None) => (model_id, "main".to_string()),
            (None, Some(revision)) => (default_model, revision),
            (None, None) => (default_model, default_revision),
        };
        let use_pth = true;
        let repo = Repo::with_revision(model_id, RepoType::Model, revision);
        let (config_filename, tokenizer_filename, weights_filename) = {
            let api = Api::new()?;
            let api = api.repo(repo);
            let config = api.get("config.json")?;
            let tokenizer = api.get("tokenizer.json")?;
            let weights = if use_pth {
                api.get("pytorch_model.bin")?
            } else {
                api.get("model.safetensors")?
            };
            (config, tokenizer, weights)
        };
        let config = std::fs::read_to_string(config_filename)?;
        let mut config: Config = serde_json::from_str(&config)?;
        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;

        // tokenizer = Tokenizer::from(tokenizer
        //     .with_padding(None)
        //     .with_truncation(None)
        //     .map_err(E::msg)?);

        let vb = if use_pth {
            VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
        } else {
            unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
        };
        let approximate_gelu = false;
        if approximate_gelu {
            config.hidden_act = HiddenAct::GeluApproximate;
        }
        let model = BertModel::load(vb, &config)?;

        // Ok((model, tokenizer))
        Ok(LanguageModel{
            model,
            tokenizer,
            device
        })
    }

    pub fn encode(&self, prompt: &str) -> Result<(Tensor,Tensor)> {

        let tokens = self.tokenizer
            .encode(prompt, true)
            .map_err(E::msg)?
            .get_ids()
            .to_vec();
        let token_ids = Tensor::new(&tokens[..],&self.device )?.unsqueeze(0)?;
        let token_type_ids = token_ids.zeros_like()?;
        let embeddings =self.model.forward(&token_ids, &token_type_ids)?;
        Ok((embeddings,token_ids))
    }
}

And this the code to load the model and handle the request

async fn create_embeddings(
    payload: web::Json<QueryPayload>,
    data: web::Data<AppState>,
) -> HttpResponse {
    let payload = &payload.query;
    let predictor = data.predictor.get().unwrap().lock().await;
    let start = Instant::now();
    let (embeddings, token_ids) = match predictor.encode(payload) {
        Ok((embeddings, token_ids)) => (embeddings, token_ids),
        Err(e) => return HttpResponse::InternalServerError().body(format!("Error: {:?}", e)),
    };
    let duration = start.elapsed();
    println!("Time taken for encoding: {:?}", duration);
    // println!("Embeddings are --> {}",embeddings);
    HttpResponse::Ok().finish()
}

and

pub fn run_embedding_server(listener: TcpListener) -> Result<Server, std::io::Error> {
    let language_model = match LanguageModel::build_model_and_tokenizer() {
        Ok(model) => model,
        Err(e) => {
            eprintln!("Failed to initialize model and tokenizer: {:?}", e);
            return Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                e.to_string(),
            ));
        }
    };
    let app_state = web::Data::new(AppState {
        predictor: OnceCell::new(),
    });
    let _ = app_state
        .predictor
        .set(Arc::new(Mutex::new(language_model)))
        .is_ok();
    let server = HttpServer::new(move || {
        App::new()
            .app_data(app_state.clone()) // Pass the application state to the app
            .route("/health_check", web::get().to(health_check))
            .route("/embeddings", web::post().to(create_embeddings))
    })
    .listen(listener)?
    .run();

    Ok(server)
}

[BerserkerMother]

I was able to replicate the result(omitted the http part). Rust is almost 2x slower than Python version. I checked the tokenizer, and they took the same time to execute, which makes sense because pythkn tokenizers is just a rust binding.
There must be something in the forward pass. Maybe the pythons backend is faster than rust?

[AbhishekBose]

Yep. Seems like it. There's something off about the BertModel implementation I guess. Will have to check

[GBathie]

Hi,

I ran into the same issue, in my case using sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2: the code is also twice as slow as python (on the CPU backend).

I did some profiling, and I believe that the culprit is candle's implementation of gelu, which is the default hidden activation function for this model. Computing the activations takes 50% (!!) of the time in BertModel::forward().

If I change the model activation function to Relu, (using config.hidden_act = HiddenAct::Relu;) this operation becomes almost instantaneous, and the code takes the same time as python.