Unspecified tensor ordering in torch.distributed.all_gather causes silent aggregation failures

Direct answer

Do not infer rank from list position. The safest deterministic pattern is to attach the source rank to each tensor before calling torch.distributed.all_gather, then recover and sort by that rank after the collective. No PyTorch version or distributed backend documents a stable rank-to-index mapping for the output list.

Confirmed facts vs. likely explanation

Confirmed: The all_gather API contract does not promise that the returned list is ordered by global rank. Treating index i as rank i is undefined behavior.

Likely explanation: NCCL and Gloo frequently return rank-ordered results in small, homogeneous tests. This masks the problem locally, but the ordering is an implementation artifact, not a guarantee, and may differ across backends, versions, or cluster topologies.

Implementation for uniform tensors

If every rank contributes the same shape and dtype, prepend a scalar rank identifier to a fixed-size payload buffer before the collective:

import torch
import torch.distributed as dist

def gather_with_rank(tensor, group=None):
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)
    # Prepend rank as first element
    tagged = torch.cat([torch.tensor([rank], dtype=tensor.dtype, device=tensor.device), tensor])
    output = [torch.empty_like(tagged) for _ in range(world_size)]
    dist.all_gather(output, tagged, group=group)
    # Recover rank and sort
    recovered = [(int(t[0].item()), t[1:]) for t in output]
    recovered.sort(key=lambda x: x[0])
    return [t for _, t in recovered]

Ensure the tag is written before the collective and not overwritten in-place afterward.

Non-uniform payloads

If tensors vary in shape or dtype, all_gather cannot be used directly. Use dist.all_gather_object to collect (rank, tensor) tuples, then sort by rank. Note that all_gather_object pickles objects and carries higher latency and bandwidth overhead than tensor collectives.

When a reduction is the goal

If the operation is sum, mean, or max, prefer dist.all_reduce. It is inherently order-independent and avoids the ordering problem entirely. If only one rank needs the gathered data, use dist.gather to that root and still attach rank metadata.

Missing diagnostic detail

Are your payloads uniform in shape and dtype across ranks? If not, the tensor-tagging path above must change to all_gather_object or padding, though the need for explicit rank identification remains unchanged.

Verification and caveats

Because this topic requires current verification, treat the following as conservative guidance: write a two-process test that gathers distinct rank-tagged tensors and asserts the recovered rank set equals {0,1}; run it across NCCL and Gloo to confirm the tag-based approach works regardless of backend. Also confirm the installed PyTorch version's API reference for torch.distributed.all_gather to ensure no recent release altered collective semantics. all_gather_into_tensor shares the same ordering caveat and does not add a rank-order guarantee.