Nx vs. Python performance for sentence-transformer encoding

@steffend it has been fixed in main here: 4e21e0467ccd5ff6a54a0115f0fe79420e089f5a

You may need to have both nx and exla pointing at that, if you have any questions, please let me know.

Also, please double check that both operations return the final data, as frameworks (both Elixir and Python) can return the output tensors without the computation fully concluding.

Finally, please double check if the SentenceTransformer is indeed padding. IIRC padding is not applied on PyTorch if you are not batching.

Yes, indeed that fixes the particular error. Thank you for looking into this!
Interestingly, the performance is still the same with 8 local devices (~117 encodes/second), though the scheduler usage in the observer looks much messier:

image1924×1424 152 KB

I’ll try. I still have much to learn in the ML space. I guess what you’re trying to say is that if the Python version does not pad the input, my short test sentence would lead to wrong results? Looking through the code I think it might pad the input (sentence-transformers/sentence_transformers/models/Transformer.py at 179b659621c680371394d507683b25ba7faa0dd8 · huggingface/sentence-transformers · GitHub), but I’m not sure if that’s really the correct piece of code.

When I find the time I will also try to compare the results of the Python and Elixir code. I have a Livebook that compute the same cosine similarities as Python using Bumblebee+Axon (no serving, as the mean pooling of the serving has some issues - Bumblebee.Text.TextEmbedding output_pool crashes · Issue #216 · elixir-nx/bumblebee · GitHub). When I have more results, I’ll update the repo and this thread.

Looks like the correct piece of code to me. So it pads to the longest input sequence (so without batching that’s no padding altogether):

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
tokenizer(["hey", "hello world"], padding=True, truncation='longest_first', return_tensors="pt", max_length=100)

#=> {'input_ids': tensor([[ 101, 4931,  102,    0], [ 101, 7592, 2088,  102]]), ...}

On the contrary we always pad do the maximum sequence length, so that we only compile once.

Keep in mind you may not want to run eight instances. When I tried this, XLA took all cores and we could not push traffic enough to the serving.

@steffend @jonatanklosko @seanmor5 I have been thinking about this and it is clear that we are more performant but forcing a certain sequence length is going to be an issue because we are always working with the worst case.

I can think of two solutions to the problem. Both are based on allowing multiple sequence lengths. For example, instead of 128, we could say 16, 32, 64, 96, and 128. If we do so, we have two options:

1. Allow multiple sequence lengths in the same batch and then pad to the highest. For example, if we get 18, 23, 42, 55, and 90 on a batch, we will pad to 96.

2. Allow multiple batch keys. In the example above, 18 and 23 go to the “32-padding batch”. 42 and 55 go to the “64-padding batch” and 90 goes to the “96-padding batch”. Each batch have their own size and individual timeouts. This means better performance but you will need to balance the batch size and batch timeout accordingly (if the timeout is high, it is more likely you will always hit the timeout).

I am thinking the batch keys approach makes the most sense but I would love to hear your thoughts.

I’ve been running some tests comparing the results more thoroughly this week and will probably post an update tomorrow. I can confirm that EXLA performs better than Python when using the full sequence length. I also started playing with CUDA on AWS, but there I still need to run some more tests.

To measure the impact of the sequence length, I adapted my serving to always tokenize twice. One time with the full sequence length and then again limited to the actual sequence length of the input. The encode/second graph looks like this for EXLA (x-axis sequence length, y-axis encodes/sec):

image1200×692 32.2 KB

This is the graph for Python (not quite fair as it goes through an extra HTTP request):

image1200×697 31 KB

And finally I’m attaching the Livebook I used to generate these graphs.

All in all, Elixir and EXLA perform well. The only thing remaining is that I could not get the CPU to be fully loaded with EXLA (the same for CUDA).

The first one seems similar to what Python does, always using the longest input sequence length, if I understood that right.
I’ve been thinking about the following: couldn’t we also allow a dynamic sequence length and just in time compile when we first get an input with a specific sequence length? Further requests should then be compiled. As the sequence length is finite, this would mean that one could either pre-compile every sequence length or “warmup” the serving.

I’ve been thinking about the following: couldn’t we also allow a dynamic sequence length and just in time compile when we first get an input with a specific sequence length? Further requests should then be compiled. As the sequence length is finite, this would mean that one could either pre-compile every sequence length or “warmup” the serving.

We can do that for sure but it means you may compile the program several times. But it is something I will consider while exploring these ideas.

Having multiple variants sounds great! Both 1. and 2. make certain tradeoffs and which is better depends on the length distribution. If longer inputs are rare, then using 2. it will hit batch timeout and we will pad with empty batch items, while we may as well put some shorter inputs there. But then note that we pad on the client as part of tokenization and it impacts all of the input tensors (input ids, attention mask), but padding to higher length means we need to pad on the server. With 2. we always know what length to pad to.

Hey,

I found this thread very helpful, but I would like to ask some additional questions / clarifications!

I am trying to speed up some text_embedding creations.
I have a data migration where we want to back-fill embeddings for existing entries. I tried various batch sizes for this and settled on 500.
I tried to create the changesets with the embedding with Enum.map and Task.async_stream. But they yield the same time. I realized that all the time is spend on generating the embedding.
I also tried to create multiple servings, both manually and with nimble_pool, but the results were the same. This leads me to believe that even though I had multiple servings the embedding creation is still sequential.
Then I found this thread and went back to a single serving and tried to tweak it.
All these finish the embedding changeset creation step in 55-57 seconds.

My vector size is 384.
I am caching the serving, is that a bad idea?

I tried the following setups:

1, In my config I have: config :nx, default_backend: EXLA.Backend
and create the serving by just calling Bumblebee.Text.TextEmbedding.text_embedding(model_info, tokenizer) as is. - I used this in all the scenarios above.

2, I tried the example form here:

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

This is slower, it takes 76 seconds.

How do I determine the batch size and the sequence_length?

Thank you