I've decided to leave this item for a follow-up PR to avoid expanding the scope of this PR too much:

• Turn util.py into a util folder with files for separate components, e.g. negative_mining.py, distributed.py, similarity_functions.py, huggingface.py, etc. It's getting too big.

With all of these losses tested (at least, checked that they perform better with gathering, ran without crashes, etc.), I'm ready to merge this. The test failures are related to the datasets v4 release, and the issues were fixed in #3445. I'll merge those into this branch first.

• Tom Aarsen

Hi, thank you for the great work on this project.

While reviewing the implementation of MultipleNegativesRankingLoss in a distributed setting, I noticed a potential issue with how the labels are computed when using all_gather on the candidates.

In particular, the code computes labels as:

if torch.distributed.is_initialized():
    rank = torch.distributed.get_rank()
    offset = rank * batch_size
labels = torch.arange(offset, offset + batch_size, device=anchors.device)

However, if candidates are constructed using torch.cat(embeddings[1:]) where [positives, negatives] are concatenated along dim=0, the actual positions of the positives in the gathered candidates tensor will not necessarily align with offset + i. This is because torch.cat first places all positives, then all negatives from each rank, so the true index of the positive may differ from what offset = rank * batch_size assumes.
I'm still new to distributed training and PyTorch, so please forgive me if I misunderstood the logic！

Not sure I get you correctly, but essentially since we are only gathering the queries, this offset is meant to align those with their corresponding positive documents, which as you pointed out are always the first n documents (because we concatenate the positive and then the negatives).
Among those n documents, the positives for a query are at [rank * batch_size: rank* batch_size + 1] (considering we concatenate following the rank order). It is just meant to skip the first positives that corresponds to the positives of previous rank.

I think @Batwu might be on to something though, although I'm not 100% sure yet. If we have e.g. 1 positive and 8 negatives with 4 devices, then the candidates that are gathered are structured a bit like PNNNNNNNNPNNNNNNNNPNNNNNNNNPNNNNNNNN.

But the offsetting with labels only offsets by batch_size, resulting in:

PNNNNNNNNPNNNNNNNNPNNNNNNNNPNNNNNNNN
1
 2
  3
   4

instead of offsetting by batch_size * (1 + num_negatives), resulting in:

PNNNNNNNNPNNNNNNNNPNNNNNNNNPNNNNNNNN
1
         2
                  3
                           4

I'll have to do some more tests, but I think it's likely that the tests I ran only used an anchor and positive, and no negatives. Otherwise I would have noticed a degradation in performance presumably.

• Tom Aarsen

That is actually totally true, I am very sorry I had the PyLate implementation in mind, where you can see that we do not concatenate the positives/negatives before gathering the elements (the "concatenation" is done at the score level), which in this case create PPPNNNNNNNN.
I believe you could simply move the concat operation to after the gathering to fix the issue no?

Again, deeply sorry @Batwu, should have checked the code of Tom instead of answering based on the code used as base.