Nx vs. Python performance for sentence-transformer encoding

Hey there,

I’ve got a project where I need to encode sentences using a sentence-transformer model. Currently, I’m using Python and the sentence-transformer package, but as the rest of the project is in Elixir I’d like to switch to Nx instead.

Using Bumblebee and Axon, I already built a small proof of concept and with the recent addition of a text embedding serving to Bumblebee, I wanted to do some quick benchmark to see how many encodes I can achieve on my CPU.

tl;dr: with a simple Python script I can achieve ~115 encodes per second with ~350% CPU load (-> ~4 Cores) on my MacBook Pro (M1 Max) and ~190 encodes per second when starting two separate Python processes (nearly full CPU utilization). Using Elixir and Nx I can only achieve ~55 encodes per second, while the average latency is more than double. Elixir also only achieves ~300% CPU usage. Starting multiple BEAM instances I can get to ~95 encodes per second with full CPU utilization.

The last point is the main one I’m interested in: there seems to be some kind of bottleneck that prevents me from achieving a similar performance to Python using only a single BEAM process. Has someone an idea why that’s the case? (It’s very possible that I’m just doing something wrong!). I expected the BEAM to be able to use all cores for encoding.

Apart from that, it seems like with full CPU utilization, I can only achieve half of the encode performance of Python using Nx, so there seem to be other factors in play too.

I’ve documented this and the code snippets here: GitHub - SteffenDE/nx-sentence-transformer-bench · GitHub

37 comments

/nx #benchmark

7 5627 37

2023-08-01 15:57:22 UTC

Marked As Solved

steffend

Phoenix Core Team

So to summarize this topic: Elixir + Nx perform quite well for generating sentence embeddings. It is very important to make sure that the model is compiled for the right sequence lengths though, as the input is padded.
For varying input lengths, the newest Bumblebee code on GitHub now supports specifying multiple sequence lengths, as seen in post 36.

Another option is using Ortex and an ONNX model. Here is an example livebook that uses this approach: Running the all-mpnet-base-v2 sentence transformer in Elixir using Ortex · GitHub. In that case, using a serving does not improve the performance that much, at least when running on CPU. Therefore one can also just call Ortex.run directly.

Thanks for all the replies and insights. I’m happy to see that this lead to some improvements in Bumblebee and Nx!

I’m also happy to report that we’re currently working on moving our production setup from Python to Elixir+Nx at my job

Post #38

Also Liked

josevalim

Creator of Elixir

I changed Nx main so we raise if the wrong option is given when starting the serving.

Post #6

seanmor5

Author of Genetic Algorithms in Elixir

Hey @steffend there is a bit of a difference between backend and compiler. You can read about it some here (it may be somewhat outdated): Nx Tip of the Week #6 - Compiler or Backend?

First I ran your benchmarks and got:

Running 1m test @ http://127.0.0.1:5001
  8 threads and 32 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   550.10ms   30.31ms 851.38ms   99.08%
    Req/Sec     6.76      2.49    30.00     96.83%
  3488 requests in 1.00m, 439.41KB read
Requests/sec:     58.03
Transfer/sec:      7.31KB

Then changing the serving to look like:

serving = Bumblebee.Text.TextEmbedding.text_embedding(model_info, tokenizer,
  compile: [batch_size: 32, sequence_length: 8],
  defn_options: [compiler: EXLA]
)

We get:

Running 1m test @ http://127.0.0.1:5001
  8 threads and 32 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency    25.69ms    2.21ms  91.61ms   95.06%
    Req/Sec   156.34     12.48   202.00     86.03%
  74832 requests in 1.00m, 9.21MB read
Requests/sec:   1245.80
Transfer/sec:    156.94KB

So pretty significant speed up just compiling the serving. There are some other config options you can mess with, but you probably won’t get much more of a speed up than that

Post #4