Nx vs. Python performance for sentence-transformer encoding

One thing I noticed is that it does not look like you are setting the compiler for your Nx serving, so you are losing a lot of optimizations there. Try setting defn_options: [compiler: EXLA] when creating the serving

I’m also not sure what the batch size you’re setting is. You can fiddle with higher and lower batch sizes to see if it improves latency.

Servings also have some built in latency, im not familiar with how the benchmark works but you can fiddle with batch timeout settings to achieve better latency as well.

Finally, if the server sends sequences of different lengths, you eat a compilation cost with every request. You should set a static sequence length

Thank you for the suggestion!

I think Nx.global_default_backend(EXLA.Backend) might already do this? At least I don’t measure any real difference when setting this on my serving. In general I think that the serving is not the limiting factor. I added a script that does not use the Nx.Serving at all, basically just calling Axon.predict and the performance is very similar (nx_axon.exs).

I also tried with different batch sizes and batch timeouts, but again without any measurable differences.

Concerning sequence lengths: that shouldn’t be an issue here as the benchmark is always encoding the same sentence, but good to know!

The main question I have is if there is some bottleneck with EXLA and the dirty NIF schedulers maybe?

image1924×1424 150 KB

Here you can see the scheduler usage while running the benchmark 3 times for 10 seconds. Looks like only one dirty cpu scheduler is used at a time, although which one changes. I’m no export on NIFs at all, so maybe that’s some common knowledge, but if Nx can only use one dirty scheduler at a time, this might become a bottleneck in other cases as well? Only speculations on my side though.

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

Oh wow, that’s indeed a very significant difference. I now also now what I did wrong: I tried to set the defn_options in the child specification instead of the serving function…

I’ll update the repo with the updated results later. Thank you!

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

To be fair to Python, I realized that this is probably just because the sequence length was limited to 8. I am pretty sure that the sequence length the python library uses is 128 (see sentence-transformers/all-MiniLM-L6-v2 · Hugging Face):

The sequence length was limited to 128 tokens.

When using a sequence length of 128, Bumblebee+EXLA achieves ~120 encodes per second, which is basically the same performance as the Python server.

This brings us back to what I was wondering in Nx vs. Python performance for sentence-transformer encoding - #3 by steffend the BEAM with Bumblebee+EXLA does not seem to be able to fully utilize all CPU cores. If I start two instances of the nx_serving script on different ports and then execute 2 instances of the benchmark, I can achieve ~190 encodes per second (with a reverse proxy it strangely gets slower).

Optimizing the Python code by using a dedicated WSGI server to run on 8 processes gunicorn -w 8 -b 0.0.0.0:5001 simple:app instead of the flask development server:

$ wrk http://127.0.0.1:5001 -t 8 -c 32 -d 60
Running 1m test @ http://127.0.0.1:5001
  8 threads and 32 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency    72.36ms   23.81ms 314.19ms   93.41%
    Req/Sec    56.71     12.13    90.00     61.91%
  26917 requests in 1.00m, 3.95MB read
Requests/sec:    447.96
Transfer/sec:     67.37KB

image982×464 21.8 KB

Two additional considerations:

1. batch_timeout and batch_size is going to impact on the latency and memory usage, so I recommend playing with those numbers if you haven’t yet. Does the Python version have anything along those lines?

2. EXLA assumes a computation will use all cores and it puts a lock around it. You can set XLA_FLAGS=--xla_force_host_platform_device_count=8 and it will start several CPU devices. You can then pass partitions: true to your Nx.Serving (in the child spec/sup tree). I am hoping this will at least allow you to use all cores within a single BEAM instance.

Yes, I already played with the batch settings and 32 seems to be a good batch size for the sequence length of 128. I did not play with the batch timeout yet, but latency is not my focus currently.

The Python version is very barebones and I don’t think that it performs any kind of batching at all. To batch I’d probably need to use something like Ray Serve: Scalable and Programmable Serving — Ray 2.56.0.

Ah that’s interesting and could very well explain what I’m seeing!

I tried this here (nx-sentence-transformer-bench/nx_serving_partitions.exs at main · SteffenDE/nx-sentence-transformer-bench · GitHub) and seeing lots of errors when benchmarking. Seems like EXLA does not like this:

17:49:30.868 [error] GenServer #PID<0.6906.0> terminating
** (stop) exited in: Nx.Serving.local_batched_run(MyServing, ["this is a test"])
    ** (EXIT) an exception was raised:
        ** (RuntimeError) Expected buffer to be placed on device 7
            (exla 0.5.3) lib/exla/executable.ex:56: EXLA.Executable.unwrap!/1
            (exla 0.5.3) lib/exla/executable.ex:19: EXLA.Executable.run/3
            (exla 0.5.3) lib/exla/defn.ex:346: EXLA.Defn.maybe_outfeed/7
            (stdlib 4.3.1.1) timer.erl:235: :timer.tc/1
            (exla 0.5.3) lib/exla/defn.ex:283: anonymous fn/7 in EXLA.Defn.__compile__/4
            (nx 0.5.3) lib/nx/defn.ex:313: anonymous fn/4 in Nx.Defn.compile/3
            (nx 0.5.3) lib/nx/serving.ex:1107: anonymous fn/2 in Nx.Serving.Default.handle_batch/3
            (nx 0.5.3) lib/nx/serving.ex:957: anonymous fn/3 in Nx.Serving.server_task_or_enqueue/3
    (nx 0.5.3) lib/nx/serving.ex:620: Nx.Serving.local_batched_run!/3
    nx_serving_partitions.exs:33: MyPlug.call/2
    (bandit 1.0.0-pre.9) lib/bandit/pipeline.ex:110: Bandit.Pipeline.call_plug/2
    (bandit 1.0.0-pre.9) lib/bandit/pipeline.ex:25: Bandit.Pipeline.run/6
    (bandit 1.0.0-pre.9) lib/bandit/http1/handler.ex:27: Bandit.HTTP1.Handler.handle_data/3
    (bandit 1.0.0-pre.9) lib/bandit/delegating_handler.ex:18: Bandit.DelegatingHandler.handle_data/3
    (bandit 1.0.0-pre.9) /Users/steffen/Library/Caches/mix/installs/elixir-1.14.5-erts-13.2.2.1/9825c67976b12b773e3fa7c81710c74f/deps/thousand_island/lib/thousand_island/handler.ex:399: Bandit.DelegatingHandler.handle_continue/2
    (stdlib 4.3.1.1) gen_server.erl:1123: :gen_server.try_dispatch/4

To reproduce:

$ XLA_FLAGS=--xla_force_host_platform_device_count=8 elixir nx_serving_partitions.exs
# then
$ wrk http://127.0.0.1:5001 -t 8 -c 32

Nice, I will investigate. Also, I recommend playing a bit with the timeout just in case (try 10ms and 1000ms as a double check).