Fix gather_for_metrics remainder when a sized sampler selects a subset

src/accelerate/data_loader.py

@@ -403,11 +403,55 @@ def begin(self):
         "Prepares the gradient state for the current dataloader"
         self.reset()
         with suppress(Exception):
-            if not self._drop_last:
-                length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            # ``DataLoader.drop_last`` stays False when a ``BatchSampler(drop_last=True)``
+            # is passed, so also honor the batch sampler's own ``drop_last``.
+            if not self._effective_drop_last():
+                length = self._length_for_remainder()
                 self.remainder = length % self.total_batch_size
         self.gradient_state._add_dataloader(self)
 
+    def _effective_drop_last(self) -> bool:
+        if getattr(self, "_drop_last", False):
+            return True
+        batch_sampler = getattr(self, "batch_sampler", None)
+        if batch_sampler is None and isinstance(getattr(self, "sampler", None), BatchSampler):
+            batch_sampler = self.sampler
+        seen = set()
+        while batch_sampler is not None and id(batch_sampler) not in seen:
+            seen.add(id(batch_sampler))
+            if getattr(batch_sampler, "drop_last", False):
+                return True
+            batch_sampler = getattr(batch_sampler, "batch_sampler", None)
+        return False
+
+    def _length_for_remainder(self) -> int:
+        """Number of samples this loader will yield, for final-batch remainder math.
+
+        Prefer a sized sampler when present (e.g. ``SubsetRandomSampler`` over a larger
+        map-style dataset). Falling back to ``len(dataset)`` alone drops or duplicates
+        samples in ``gather_for_metrics`` when the sampler selects a different count
+        (#4410). Iterable datasets can still expose ``total_dataset_length``.
+
+        ``batch_sampler.sampler`` is checked before ``self.sampler``: after
+        ``Accelerator.prepare``, the top-level ``sampler`` can be a default
+        ``SequentialSampler`` over the full dataset while the real selection lives
+        on ``batch_sampler.sampler``.
+
+        Callers should invoke ``set_epoch`` before ``begin`` so epoch-dependent
+        sampler lengths are current when remainder is computed.
+        """
+        sampler_candidates = (
+            getattr(getattr(self, "batch_sampler", None), "sampler", None),
+            getattr(getattr(getattr(self, "batch_sampler", None), "batch_sampler", None), "sampler", None),
+            getattr(self, "sampler", None),
+        )
+        for sampler in sampler_candidates:
+            if sampler is None or isinstance(sampler, BatchSampler):
+                continue
+            with suppress(TypeError, NotImplementedError, AttributeError):
+                return len(sampler)
+        return getattr(self.dataset, "total_dataset_length", len(self.dataset))
+
     def end(self):
         "Cleans up the gradient state after exiting the dataloader"
         self.gradient_state._remove_dataloader(self)
@@ -578,9 +622,9 @@ def adjust_state_dict_for_prefetch(self):
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.synchronized_generator)
-        self.begin()
-
+        # set_epoch before begin so remainder uses the current epoch's sampler length
         self.set_epoch(self.iteration)
+        self.begin()
         dataloader_iter = self.base_dataloader.__iter__()
         # We iterate one batch ahead to check when we are at the end
         try:
@@ -873,8 +917,9 @@ def _fetch_batches(self, iterator):
         return batch, batch_info
 
     def __iter__(self):
-        self.begin()
+        # set_epoch before begin so remainder uses the current epoch's sampler length
         self.set_epoch(self.iteration)
+        self.begin()
         main_iterator = None
         if is_torch_version(">=", "2.0.1"):
             # NOTE PyTorch DataLoader adds forward compatibilities for DataPipes, which broadcasts

tests/test_data_loader.py

@@ -18,7 +18,7 @@
 import pytest
 import torch
 from parameterized import parameterized
-from torch.utils.data import BatchSampler, DataLoader, IterableDataset
+from torch.utils.data import BatchSampler, DataLoader, IterableDataset, Sampler, Subset, SubsetRandomSampler
 
 from accelerate import Accelerator, PartialState
 from accelerate.data_loader import (
@@ -631,6 +631,87 @@ def test_skip_first_batches_preserves_drop_last(self):
         assert [batch.tolist() for batch in resumed] == [[4, 5]]
         assert resumed.remainder == -1
 
+    def test_remainder_uses_sampler_length_for_gather_for_metrics(self):
+        # #4410: SubsetRandomSampler selects 10 of 21; remainder must use 10, not len(dataset)=21
+        acc = Accelerator(cpu=True)
+        ds = range(21)
+        sampler = SubsetRandomSampler(range(10), generator=torch.Generator().manual_seed(42))
+        loader = acc.prepare(DataLoader(ds, batch_size=4, sampler=sampler))
+
+        got = []
+        for batch in loader:
+            got.extend(acc.gather_for_metrics(batch).tolist())
+
+        assert loader.remainder == 10 % loader.total_batch_size
+        assert sorted(got) == list(range(10))
+
+        # Control: Subset with the same sampler length keeps agreeing
+        control_ds = Subset(range(21), range(10))
+        control_sampler = SubsetRandomSampler(range(10), generator=torch.Generator().manual_seed(42))
+        control_loader = acc.prepare(DataLoader(control_ds, batch_size=4, sampler=control_sampler))
+        control_got = []
+        for batch in control_loader:
+            control_got.extend(acc.gather_for_metrics(batch).tolist())
+        assert sorted(control_got) == list(range(10))
+        assert sorted(got) == sorted(control_got)
+        acc.end_training()
+
+    def test_remainder_respects_batch_sampler_drop_last(self):
+        # BatchSampler(drop_last=True) does not set DataLoader.drop_last; remainder must stay unused
+        # so gather_for_metrics keeps all yielded samples (8 of 10), not truncate to 6.
+        acc = Accelerator(cpu=True)
+        ds = range(21)
+        sampler = SubsetRandomSampler(range(10), generator=torch.Generator().manual_seed(0))
+        batch_sampler = BatchSampler(sampler, batch_size=4, drop_last=True)
+        loader = acc.prepare(DataLoader(ds, batch_sampler=batch_sampler))
+
+        got = []
+        for batch in loader:
+            got.extend(acc.gather_for_metrics(batch).tolist())
+
+        assert loader.remainder == -1
+        assert len(got) == 8
+        assert sorted(got) == sorted(list(sampler)[:8]) or len(got) == 8
+        # All 8 samples come from the 10-index subset
+        assert set(got).issubset(set(range(10)))
+        acc.end_training()
+
+    def test_remainder_uses_length_after_set_epoch(self):
+        # Sampler length can change across epochs; begin() must run after set_epoch.
+        class GrowingSampler(Sampler):
+            def __init__(self, n0=5):
+                self.epoch = 0
+                self.n0 = n0
+
+            def set_epoch(self, epoch):
+                self.epoch = epoch
+
+            def __iter__(self):
+                return iter(range(len(self)))
+
+            def __len__(self):
+                return self.n0 + self.epoch
+
+        acc = Accelerator(cpu=True)
+        ds = range(21)
+        sampler = GrowingSampler(n0=5)
+        loader = acc.prepare(DataLoader(ds, batch_size=4, sampler=sampler))
+
+        # epoch 0 via iteration=0 -> length 5, remainder 5 % 4 = 1
+        got0 = []
+        for batch in loader:
+            got0.extend(acc.gather_for_metrics(batch).tolist())
+        assert sorted(got0) == list(range(5))
+        assert loader.remainder == 1
+
+        # second pass: iteration advanced to 1 -> length 6, remainder 6 % 4 = 2
+        got1 = []
+        for batch in loader:
+            got1.extend(acc.gather_for_metrics(batch).tolist())
+        assert sorted(got1) == list(range(6))
+        assert loader.remainder == 2
+        acc.end_training()
+
     def test_end_of_dataloader(self):
         dataloader = DataLoaderShard(list(range(16)), batch_size=4)
         for idx, _ in enumerate(dataloader):