distributed

sentence_transformers.util.distributed.all_gather(tensor: Tensor, with_grad: bool = False) → Tensor

Gathers a tensor from each distributed rank into a list. Always retains gradients for the local rank’s tensor, and optionally retains gradients for the gathered tensors if with_grad is True.

Parameters:

tensor (torch.Tensor) – The tensor to gather from each rank.
with_grad (bool, optional) – If True, the local rank’s tensor retains its gradients. Defaults to False.

Returns:

A tensor containing the gathered tensors from all ranks, concatenated along the first dimension. If torch.distributed is not available or not initialized, returns the original tensor.

Return type:

torch.Tensor

sentence_transformers.util.distributed.all_gather_padded(tensor: Tensor, mask: Tensor, with_grad: bool = False) → tuple[Tensor, Tensor]

All-gather a (B, T, D) token-embedding tensor and its (B, T) mask across ranks, padding the token axis to the cross-rank max T first.

all_gather requires every rank to contribute an identically-shaped tensor, but multi-vector batches pad each column to its own per-rank batch-longest T, which differs across ranks. This reduces the global max T, pads both the embeddings and the mask up to it (keeping them aligned), then gathers. Embeddings are gathered with with_grad. The mask never carries a gradient.

Parameters:

tensor (torch.Tensor) – (B, T, D) token embeddings to gather.
mask (torch.Tensor) – (B, T) mask to gather, padded with False to match tensor.
with_grad (bool, optional) – Retain gradients for the embeddings (see all_gather()).

Returns:

Gathered (sum(B), T_max, D) embeddings and (sum(B), T_max) mask. Without an initialised process group this only forwards to all_gather() (no padding needed, all_gather returns the local tensor).

Return type:

tuple[torch.Tensor, torch.Tensor]

sentence_transformers.util.distributed.all_gather_with_grad(tensor: Tensor) → Tensor

Gathers a tensor from each distributed rank into a list, retaining gradients for the local rank’s tensor.

Parameters:: tensor (torch.Tensor) – The tensor to gather from each rank.
Returns:: A tensor containing the gathered tensors from all ranks, concatenated along the first dimension. If torch.distributed is not available or not initialized, returns the original tensor.
Return type:: torch.Tensor

sentence_transformers.util.distributed.get_rank() → int: The rank of the current process in the distributed group, or 0 when not distributed.

sentence_transformers.util.distributed.get_world_size() → int: The number of processes in the distributed group, or 1 when not distributed.