Multi-GPU for more tasks

[Muennighoff]

It'd be great if we could figure out using multiple gpus on tasks other than BEIR.
E.g. RedditClusteringP2P takes >20h for a 5.8B model with embeddings of 4096 dimensions.

[dylan-hum]

I wrote a wrapper for the model like this, it seems to be working so far.

class MultiGPUModel(object):
    def __init__(self, sentence_transformer: SentenceTransformer):
        self.transformer = sentence_transformer
        self.gpu_pool = self.transformer.start_multi_process_pool()

    def encode(self, sentences, **kwargs):
        return self.transformer.encode_multi_process(sentences, self.gpu_pool, **kwargs)

Then you can pass it like this:

from mteb import MTEB
from sentence_transformers import SentenceTransformer

# Define the sentence-transformers model name
model_name = "average_word_embeddings_komninos"

model = SentenceTransformer(model_name)
multigpu_model = MultiGPUModel(model)
evaluation = MTEB(tasks=["Banking77Classification"])
results = evaluation.run(multigpu_model , output_folder=f"results/{model_name}")

[Muennighoff]

I wrote a wrapper for the model like this, it seems to be working so far.

class MultiGPUModel(object):
    def __init__(self, sentence_transformer: SentenceTransformer):
        self.transformer = sentence_transformer
        self.gpu_pool = self.transformer.start_multi_process_pool()

    def encode(self, sentences, **kwargs):
        return self.transformer.encode_multi_process(sentences, self.gpu_pool, **kwargs)

Then you can pass it like this:

from mteb import MTEB
from sentence_transformers import SentenceTransformer

# Define the sentence-transformers model name
model_name = "average_word_embeddings_komninos"

model = SentenceTransformer(model_name)
multigpu_model = MultiGPUModel(model)
evaluation = MTEB(tasks=["Banking77Classification"])
results = evaluation.run(multigpu_model , output_folder=f"results/{model_name}")

As highlighted in this wrapper, I think this can be done much more easily on the user side and thus does not need to be in MTEB ; Also see #233 :)