Sentence-transformers: Multi GPU utilization not optimal.

[PhilipMay]

Hi,
I am using an AWS instance with 4 * Tesla V100 GPU and 32 vCPU.

Code is this:

    model = SentenceTransformer('../models/sentence_bert_02')
    pool = model.start_multi_process_pool(encode_batch_size=2000)
    embeddings = model.encode_multi_process(texts, pool)
    model.stop_multi_process_pool(pool)

Here is an typical nvidia-smi output:

| NVIDIA-SMI 450.51.05    Driver Version: 450.51.05    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:00:1B.0 Off |                    0 |
| N/A   68C    P0   262W / 300W |  11748MiB / 16160MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  On   | 00000000:00:1C.0 Off |                    0 |
| N/A   52C    P0    63W / 300W |  12362MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  Tesla V100-SXM2...  On   | 00000000:00:1D.0 Off |                    0 |
| N/A   55C    P0    75W / 300W |  13726MiB / 16160MiB |      3%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  Tesla V100-SXM2...  On   | 00000000:00:1E.0 Off |                    0 |
| N/A   57C    P0    70W / 300W |  12362MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     52400      C   ...conda3/envs/hf/bin/python    11745MiB |
|    1   N/A  N/A     52473      C   ...conda3/envs/hf/bin/python    12359MiB |
|    2   N/A  N/A     52550      C   ...conda3/envs/hf/bin/python    13723MiB |
|    3   N/A  N/A     52620      C   ...conda3/envs/hf/bin/python    12359MiB |
+-----------------------------------------------------------------------------+

The batch size seems to be good since the GPU memory is almost full.
But the GPUs are often at 0% or almost at 0%.

I just wanted to report back because it seems that there is a bottlenack. Is it possible to
specify the multiprocessing pool size? Can I use more than 4 CPUs for tokenization?

PS: len(texts) is > 4_000_000 btw.

[PhilipMay]

Two or 3 GPUs of the 4 GPU total are always idle.

[nreimers]

Not sure why this happens. Just tested it on a DGX-2 machine with V100 GPUs and it was using all 4 GPUs with near 100% utilization. As model I used a bert-base model

The logic is the following for this method:

• It chunks the sentences into smaller packages. By default, into 5000 sentences.
• These sentences are tokenized
• Then added to a Queue from where the different processes fetch the data
• Each process fetches data, computes the embeddings, and returns it to the parent process
• The parent process merges all the different chunks back and returns the final list of embeddings

Bottlenecks can be:
Tokenization. Computing the embedding might be faster than the tokenizer. Then only one process will be used. You could try to pre-tokenize the inputs and then set is_pretokenized=True:

sentences = [model.tokenize(sent) for sentences]
embeddings = model.encode_multi_process(sentences, pool, is_pretokenized=True)

You could also try to use the Fast BERT tokenizer from HF and see if it makes a difference.

Another bottle neck might be the inter-process communication, that it takes too long to communicate the data (sentences as input, embeddings as output).

[PhilipMay]

Maybe it is an AWS / cloud specific issue. I will try the pretokenization.

One idea: Would it be possible to increase (specify) the number of parallel workers that do the tokenization?

[nreimers]

Multi-process tokenization was removed, as it caused to many issues.

But as you can pass pre-tokenized sentences, you can create your own pool for the tokenization, similar like this:

from torch.multiprocessing import Pool

if __name__ == '__main__':
    with Pool(5) as p:
        tokenized_sentences = list(p.imap(model.tokenize, sentences, chunksize=1000))

In case you use BERT, also have a look at the fast tokenizers from HF, which apply multi-thread / multi-processing as far as I know.

Let me know if the issue is solved after switching to pre-tokenized inputs.

[PhilipMay]

Might be an AWS specific issue. Closing this. Thanks.