Length grouping for iterable datasets
Author: ssharpe42
Pull request: no
Labels: name: Feature request

### Feature request

Similar to non-iterable datasets I would like functionality to group batches by similar length inputs. 

Example of how this could be implemented is using a buffer to preload batches and sort in IterableDatasetShard in `accelerate` using the defined column for length grouping. 

```
class IterableDatasetShard(IterableDataset):
    def __init__(
        self,
        dataset: IterableDataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
        split_batches: bool = False,
        group_by_length_buffer: int = 1, ## how many batches to preload
        length_column_name: str = "", ## which column determines length
    ):
        if split_batches and batch_size > 1 and batch_size % num_processes != 0:
            raise ValueError(
                f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) "
                f"needs to be a round multiple of the number of processes ({num_processes})."
            )
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index
        self.split_batches = split_batches
        self.group_by_length_buffer = group_by_length_buffer
        self.length_column_name = length_column_name

    def set_epoch(self, epoch):
        self.epoch = epoch
        if hasattr(self.dataset, "set_epoch"):
            self.dataset.set_epoch(epoch)

    def __len__(self):
        # We will just raise the downstream error if the underlying dataset is not sized
        if self.drop_last:
            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
        else:
            return (
                math.ceil(len(self.dataset) / (self.batch_size * self.num_processes))
                * self.batch_size
            )

    def __iter__(self):
        if (
            not hasattr(self.dataset, "set_epoch")
            and hasattr(self.dataset, "generator")
            and isinstance(self.dataset.generator, torch.Generator)
        ):
            self.dataset.generator.manual_seed(self.epoch)
        real_batch_size = (
            self.batch_size if self.split_batches else (self.batch_size * self.num_processes)
        ) * self.group_by_length_buffer
        process_batch_size = (
            (self.batch_size // self.num_processes) if self.split_batches else self.batch_size
        ) * self.group_by_length_buffer
        process_slice = range(
            self.process_index * process_batch_size, (self.process_index + 1) * process_batch_size
        )

        first_batch = None
        current_batch = []
        for element in self.dataset:
            current_batch.append(element)
            # Wait to have a full batch before yielding elements.
            if len(current_batch) == real_batch_size:
                process_current_batch = [current_batch[i] for i in process_slice]
                process_current_batch.sort(key=lambda x: len(x[self.length_column_name]))
                yield from process_current_batch
                if first_batch is None:
                    first_batch = current_batch.copy()
                current_batch = []

        # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning.
        if not self.drop_last and len(current_batch) > 0:
            if first_batch is None:
                first_batch = current_batch.copy()
            while len(current_batch) < real_batch_size:
                current_batch += first_batch
            for i in process_slice:
                yield current_batch[i]


```





### Motivation

This will help improve efficiency for our datasets with very large variance in sequence lengths. 

### Your contribution

I could submit an initial PR .

Comments:
@ArthurZucker, @amyeroberts

Sure! If you have a PR it is most welcome 🤗 let's also add in some doc about why people should us it and 🚀 

Hey @ssharpe42, sorry for the very slow reply.

Still valid `group_by_length` no-ops for iterable datasets because the sampler bails when the dataset has no length https://github.com/huggingface/transformers/blob/9f909e0f7027285d83bb6addc88155f1b80244ab/src/transformers/trainer.py#L1016-L1017
But it's been two years with no PR and no other askers, so I'm closing to keep the backlog honest. Not a rejection: @ArthurZucker welcomed a PR in 2024 and that stands; open one (even a draft) and we'll review it and reopen this. FWIW IterableDatasetShard now lives in transformers so a prototype can go here rather than in accelerate.

The silent no-op is a separate bug, it used to error (#15437), dropped in #42265. Tracking that on its own.