[Improvement] Update DistributedLengthGroupedSampler to allow customizing length function

src/transformers/trainer_pt_utils.py

@@ -24,7 +24,7 @@
 import re
 import sys
 import warnings
-from collections.abc import Iterator, Mapping
+from collections.abc import Callable, Iterator, Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from itertools import chain
@@ -508,7 +508,50 @@ def __init__(
         drop_last: bool = False,
         lengths: list[int] | None = None,
         model_input_name: str | None = None,
+        length_func: Callable[[Any], int] | None = None,
+        mega_batch_mult: int | None = None,
     ):
+        """
+        **Warning**: This sampler may be slow to initialize if lengths is not provided. It is
+        recommended to cache the resulting lengths object after the first run to speed up subsequent runs.
+
+        Args:
+            batch_size (`int`):
+                The batch size to use for sampling.
+            dataset (`Dataset`, *optional*):
+                The dataset to sample from. Either `dataset` or `lengths` must be provided.
+            num_replicas (`int`, *optional*):
+                Number of processes participating in distributed training. Will default to
+                the dist.get_world_size() if not provided.
+            rank (`int`, *optional*):
+                Rank of the current process within `num_replicas`. Will default to the
+                dist.get_rank() if not provided.
+            seed (`int`, *optional*, defaults to `0`):
+                Random seed used for shuffling.
+            drop_last (`bool`, *optional*, defaults to `False`):
+                If `True`, the sampler will drop the tail of the data to make it evenly
+                divisible across the number of replicas.
+            lengths (`list[int]`, *optional*):
+                Pre-computed lengths of the dataset items. If not provided, lengths will be
+                inferred from the dataset using `length_func` or the default method.
+            model_input_name (`str`, *optional*):
+                The name of the key in the dataset items to use for computing lengths.
+                Defaults to `"input_ids"` if not specified. Ignored if length_func is provided.
+            length_func (`Callable[[Any], int]`, *optional*):
+                A function that takes a dataset item and returns its length. If not provided,
+                the length will be inferred from the `model_input_name` key.
+            mega_batch_mult (`int`, *optional*):
+                The sampler takes mega_batch_mult * batch_size number of samples into memory
+                before sorting them by length. This parameter controls the size of these mega
+                batches. If not provided, it will default to min(len(dataset) // (batch_size * 4), 50).
+
+        Raises:
+            ValueError: If neither `dataset` nor `lengths` is provided.
+            RuntimeError: If distributed package is not available when `num_replicas` or `rank`
+                is not provided.
+            ValueError: If lengths cannot be automatically inferred from the dataset.
+        """
+
         if dataset is None and lengths is None:
             raise ValueError("One of dataset and lengths must be provided.")
         if num_replicas is None:
@@ -525,15 +568,23 @@ def __init__(
         self.rank = rank
         self.epoch = 0
         self.drop_last = drop_last
+        self.mega_batch_mult = mega_batch_mult
 
         if lengths is None:
-            model_input_name = model_input_name if model_input_name is not None else "input_ids"
-            if not isinstance(dataset[0], (dict, BatchEncoding)) or model_input_name not in dataset[0]:
-                raise ValueError(
-                    "Can only automatically infer lengths for datasets whose items are dictionaries with an "
-                    f"'{model_input_name}' key."
-                )
-            lengths = [len(feature[model_input_name]) for feature in dataset]
+            if length_func is None:
+                model_input_name = model_input_name if model_input_name is not None else "input_ids"
+                if not isinstance(dataset[0], (dict, BatchEncoding)) or model_input_name not in dataset[0]:
+                    raise ValueError(
+                        "Can only automatically infer lengths for datasets whose items are dictionaries with an "
+                        f"'{model_input_name}' key."
+                    )
+
+                def _length_func(x):
+                    return len(x[model_input_name])
+
+                length_func = _length_func
+
+            lengths = [length_func(feature) for feature in dataset]
         elif isinstance(lengths, torch.Tensor):
             logger.info(
                 "If lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to"
@@ -559,7 +610,9 @@ def __iter__(self) -> Iterator:
         # Deterministically shuffle based on epoch and seed
         g = torch.Generator()
         g.manual_seed(self.seed + self.epoch)
-        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g)
+        indices = get_length_grouped_indices(
+            self.lengths, self.batch_size, mega_batch_mult=self.mega_batch_mult, generator=g
+        )
 
         if not self.drop_last:
             # add extra samples to make it evenly divisible

tests/trainer/test_trainer_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import copy
+import random
 import unittest
 import warnings
 
@@ -28,7 +29,7 @@
 if is_torch_available():
     import torch
     from torch import nn
-    from torch.utils.data import IterableDataset
+    from torch.utils.data import BatchSampler, Dataset, IterableDataset
 
     from transformers.modeling_outputs import SequenceClassifierOutput
     from transformers.tokenization_utils_base import BatchEncoding
@@ -75,6 +76,19 @@ def __iter__(self):
                 number = torch.rand(1, generator=self.generator).item()
                 stop = number < self.p_stop
 
+    class TensorListDataset(Dataset[torch.Tensor]):
+        tensors: list[torch.Tensor]
+
+        def __init__(self, tensors: list[torch.Tensor]) -> None:
+            assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors), "Size mismatch between tensors"
+            self.tensors = tensors
+
+        def __getitem__(self, index):
+            return self.tensors[index]
+
+        def __len__(self):
+            return len(self.tensors)
+
 
 @require_torch
 class TrainerUtilsTest(unittest.TestCase):
@@ -161,6 +175,70 @@ def test_distributed_length_grouped(self):
         # The indices should be a permutation of range(100)
         self.assertEqual(sorted(indices_process_0 + indices_process_1), list(range(100)))
 
+    def test_distributed_length_grouped_sampler(self):
+        # Simulate a dataset with dict items containing input_ids
+        data = []
+        for length in range(10, 110, 10):  # 10, 20, 30, ..., 100
+            for _ in range(10):
+                data.append({"input_ids": torch.randn(length)})
+        random.shuffle(data)
+
+        sampler = DistributedLengthGroupedSampler(
+            batch_size=10,
+            dataset=data,
+            num_replicas=1,
+            rank=0,
+            mega_batch_mult=100,
+        )
+
+        batch_sampler = BatchSampler(sampler, batch_size=10, drop_last=False)
+        batches = list(batch_sampler)
+
+        next_batch = batches[0]
+        assert len(next_batch) == 10
+        assert all(len(data[i]["input_ids"]) == len(data[next_batch[0]]["input_ids"]) for i in next_batch)
+
+        other_batch = batches[1]
+        assert len(other_batch) == 10
+        assert all(len(data[i]["input_ids"]) == len(data[other_batch[0]]["input_ids"]) for i in other_batch)
+
+        assert len(data[next_batch[0]]["input_ids"]) != len(data[other_batch[0]]["input_ids"])
+
+    def test_distributed_length_grouped_sampler_custom_lengths(self):
+        # Simulate a dataset where each sample has shape (1, seq_len) with random lengths
+        data = []
+        for length in range(10, 110, 10):  # 10, 20, 30, ..., 100
+            for _ in range(10):
+                data.append(torch.randn(1, length))
+        random.shuffle(data)
+
+        def length_func(sample):
+            return sample.shape[1]
+
+        sampler = DistributedLengthGroupedSampler(
+            batch_size=10,
+            dataset=TensorListDataset(data),
+            num_replicas=1,
+            rank=0,
+            length_func=length_func,
+            mega_batch_mult=100,  # Ensure entire dataset is considered for grouping
+        )
+
+        batch_sampler = BatchSampler(sampler, batch_size=10, drop_last=False)
+        batches = list(batch_sampler)
+
+        next_batch = batches[0]
+
+        assert len(next_batch) == 10
+        assert all(data[i].shape[1] == data[next_batch[0]].shape[1] for i in next_batch)
+
+        other_batch = batches[1]
+        assert len(other_batch) == 10
+        assert all(data[i].shape[1] == data[other_batch[0]].shape[1] for i in other_batch)
+
+        # Other batch should have different sequence length
+        assert data[next_batch[0]].shape[1] != data[other_batch[0]].shape[1]
+
     def test_get_parameter_names(self):
         model = nn.Sequential(TstLayer(128), nn.ModuleList([TstLayer(128), TstLayer(128)]))
         # fmt: off