Add Support for Multi-GPU Training with PyTorch Lightning

[OrjwanZaafarani]

This pull request introduces support for multi-GPU training in the Sentence Transformers library using PyTorch Lightning.

The following changes have been made:

• Updated README.md to include instructions on how to perform multi-GPU training.
• Added a new module, SentenceTransformerMultiGPU.py, to enable multi-GPU training.
• Included PyTorch Lightning in the requirements.txt file to ensure compatibility.

This enhancement will allow users to leverage the power of multiple GPUs for faster and more efficient training with Sentence Transformers.

Closes the following issues:

• Multi-GPU training #311
• Multi-GPU training #556
• How to use multi GPU in training #59
• Multi-gpu fine-tunning of sbert with roberta-base #1898
• SentenceBert Fit-Method: Capture loss, Gradient Accumulation and Multi-GPU #1340
• MultiGPU training not avalilable on Sentence Transformer #1429

[OrjwanZaafarani]

Multi-GPU Training

The following snippet of code wraps SentenceTransformer to enable multi-GPU training using Pytorch Lightning:

from SentenceTransformer_MultiGPU import SentenceTransformerMultiGPU
from sentence_transformers import losses
import lightning.pytorch as pl
model = SentenceTransformerMultiGPU("author/any_huggingface_model", losses.CosineSimilarityLoss)
trainer = pl.Trainer(max_epochs = 2,
                     accelerator = 'gpu',
                     devices = 3,
                     strategy = 'ddp_find_unused_parameters_true',
                     log_every_n_steps = 50)
trainer.fit(model = model, train_dataloader = train_dataloader, val_dataloader = val_dataloader)

Please note that the dataloader expects an InputExample object out of the dataset class. For more details, take a look at the main module SentenceTransformer_MultiGPU.py and the example examples/multigpi_training.py.

[tomaarsen]

Hello!

Thanks for creating this PR. I am intending to introduce Multi-GPU training via #2449 instead, by relying on the transformers Trainer. The idea is to fully replace the current fit training approach with that Trainer-based approach, while preserving backwards compatibility (i.e. the fit will still work, it just uses the Trainer behind the scenes).

I've kept this open in case I ran into concrete issues with #2449, but I think I would rather move forward with that PR instead.

• Tom Aarsen

[tomaarsen]

Hello!

Tthis behaviour has now been implemented in the v3.0 update. See https://sbert.net/docs/sentence_transformer/training/distributed.html for more details.

• Tom Aarsen