distributed
- sentence_transformers.util.distributed.all_gather(tensor: Tensor, with_grad: bool = False) Tensor[source]
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:
- sentence_transformers.util.distributed.all_gather_with_grad(tensor: Tensor) Tensor[source]
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: