Output Contains Duplicate Keys When Using DistributedSampler in Distributed Inference

[yfyeung]

When using s3tokenizer with PyTorch's DistributedSampler in a distributed inference setup, the output files contain duplicate keys, and the total number of keys is always a multiple of world_size.

This issue arises because DistributedSampler in PyTorch, by default, pads the dataset by repeating samples if drop_last=False, to ensure the total dataset size is divisible by world_size. This behavior is implemented in the DistributedSampler source code, where extra samples are added if len(indices) % num_replicas != 0, causing certain samples to appear multiple times across different ranks.

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[
                    :padding_size
                ]

[xingchensong]

welcome for pr