Gather tensors of unequal shapes

Description & Motivation

It would be helpful to have a function that can gather tensors across processes of unequal shapes. For example, we have two arrays of shapes: (1, 3) on rank 0 and (4, 3) on rank 1. Current all_gather will return (2, 1, 3) for rank 0 and (2, 4, 3) for rank 1. It would be nice to concatenate across a specified dimension to get (5, 3) tensor. here is my implementation:

import torch
from lightning import Fabric


def cat_across_processes(data: torch.Tensor, fabric, dim=0):
    if fabric.world_size <= 1:
        return data

    dim_len = data.shape[dim]
    all_lens = fabric.all_gather(dim_len).flatten()
    max_len = max(all_lens).item()

    # pad data
    if dim_len < max_len:
        shape = list(data.shape)
        shape[dim] = max_len - dim_len
        padding = torch.empty(shape, dtype=data.dtype, device=data.device)
        data = torch.cat([data, padding], dim=dim)

    # all gather across all processes, generates a tensor of shape (world_size, ...)
    data = fabric.all_gather(data)

    # delete padded elements and concatenate
    return torch.cat([d.narrow(dim, 0, l) for d, l in zip(data.unbind(0), all_lens)], dim=dim)


fabric = Fabric(devices=2)
fabric.launch()

if fabric.global_rank == 0:
    tensor = torch.ones(1, 3, device=fabric.device)
else:
    tensor = torch.zeros(4, 3, device=fabric.device)


gathered = cat_across_processes(tensor, fabric)

print(f"rank {fabric.global_rank}: {tensor.shape} {gathered.shape}")

if fabric.global_rank == 0:
    print(gathered)

Pitch

No response

Alternatives

No response

Additional context

No response

cc @borda @awaelchli @carmocca @justusschock

Start by tracing Fabric.all_gather and the distributed collective path it uses. Reproduce the example with tensors shaped (1, 3) and (4, 3), then determine where an API for unequal shapes and a configurable concatenation dimension belongs. Done means tensors with different lengths on the selected dimension are gathered and concatenated consistently across processes, with coverage for the documented example.