Thanks for pointing out this issue!
If you have some sort of identifier for your input (i.e. input_ids as you mentioned) can you build a mapping of {input_id: prediction} on each distributed rank? After inference on N inputs, this will be a dict of size N on each rank and you can use all_gather_object (torch.distributed.distributed_c10d — PyTorch 1.9.0 documentation) to get the entire input → prediction mapping across the world. Would something like this satisfy your use case?