Cached losses (GradCache) trigger one DDP all-reduce per mini-batch per column instead of one per step

tl;dr GradCache triggers all-reduce for every mini-batch in DDP instead of for each intended optimizer step, which leads to DDP overhead as well as hang potential, leading to a desync between DDP processes.

CachedLossMixin is the parent pf every Cached* loss. Its backward hook re-embeds and backpropagates one mini-batch at a time:

https://github.com/huggingface/sentence-transformers/blob/94d68cf9cca6498d54e4c0d3f84532cbd33a7a9f/sentence_transformers/base/losses/gradcache.py#L297-L318

So surrogate.backward() runs once per mini-batch, AND once per sentence-feature column (anchor, positive, ...). The problem is that model.no_sync() is never called. When the loss_obj.model is wrapped in DistributedDataParallel, every one of these calls is a forward+backward through the DDP module, so DDP's bucket-ready hooks fire and dispatch a real gradient all-reduce on every single mini-batch and not just once per optimizer step, which is what we would expect to happen.

For example, with batch_size=1024, mini_batch_size=32 and a 2-column loss (like anchor+positive), that's (1024/32) x 2 = 64 all-reduce dispatches for one training step, where one should have been sufficient. That may seem like just some extra overhead but, as I've experienced, it comes also with a hang risk that bugs sentence. Because there are no guards in between those 64 sequential reductions, , it is possible that one rank's all-reduce dispatch count gets out of sync with another. In that case, ProcessGroupNCCL will eventually timeout because that one rank is waiting on a collective the other one never got to.

(I hit this problem exactly when rank 0's watchdog fired on ALLREDUCE with last enqueued work: 19021 vs last completed work: 19015 (6 more collectives enqueued than the other rank ever managed to reach), and the failed collective's stack trace pointed straight at that gradcache.py line (318) above.)

We can reproduce this even without multi-gpu by just checking how the communication hooks work in a gloo group. (sentence-transformers 6.0.1, torch 2.13.0).

import torch, torch.distributed as dist, torch.nn as nn, torch.nn.functional as F
from sentence_transformers.base.losses.gradcache import CachedLossMixin

BATCH_SIZE, MINI_BATCH_SIZE = 16, 4

class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(4, 4)
    def forward(self, features):
        return {"sentence_embedding": self.linear(features["input_ids"])}

class TinyCachedLoss(nn.Module, CachedLossMixin):
    def __init__(self, model, mini_batch_size):
        super().__init__()
        self.model = model
        self.mini_batch_size = mini_batch_size
    def calculate_loss(self, reps, labels=None, *, with_backward=False):
        anchors, positives = torch.cat(reps[0]), torch.cat(reps[1])
        loss = F.cross_entropy(anchors @ positives.T, torch.arange(anchors.size(0)))
        if with_backward:
            loss.backward()
        return loss.detach()
    def forward(self, sentence_features, labels=None):
        return self.forward_cached(sentence_features, labels)

dist.init_process_group("gloo", rank=0, world_size=1, init_method="tcp://127.0.0.1:29511")
ddp_model = nn.parallel.DistributedDataParallel(TinyModel())

sync_count = 0
def counting_hook(_state, bucket) -> torch.futures.Future[torch.Tensor]:
    global sync_count
    sync_count += 1
    fut = torch.futures.Future()
    fut.set_result(bucket.buffer())
    return fut
ddp_model.register_comm_hook(state=None, hook=counting_hook)

loss_fn = TinyCachedLoss(ddp_model, mini_batch_size=MINI_BATCH_SIZE)
anchors = {"input_ids": torch.randn(BATCH_SIZE, 4)}
positives = {"input_ids": torch.randn(BATCH_SIZE, 4)}
loss_fn([anchors, positives]).backward()

print(f"mini-batches per column: {BATCH_SIZE // MINI_BATCH_SIZE}, columns: 2")
print(f"DDP all-reduce dispatched: {sync_count} times (expected: 1)")

Observed output: mini-batches per column: 4, columns: 2 DDP all-reduce dispatched: 8 times (expected: 1)

The way to fix this, I think, is to wrap every surrogate.backward() call execpt the last one (i.e. the last one across all columns and mini-batches within this step) in a model.no_sync, which would accumulate .grad and then finally flushes it in the final unguarded call.

Loss calculation determinism is the same as the current implementation but should cause less overhead in DDP and, hopefully, decrease the risk of desyncing.

I can draft a PR if you'd like.