gather_for_metrics keeps padding duplicates for a sized IterableDataset with dispatch_batches=False #4359

System Info

- `Accelerate` version: 1.16.0.dev0 (main @ aebbcc7)
- Platform: macOS-26.5-arm64-arm-64bit
- Python version: 3.12.13
- Numpy version: 2.5.3
- PyTorch version: 2.14.0
- PyTorch accelerator: MPS (repro below runs on CPU, 2 processes via torchrun)
- `Accelerate` default config: Not found

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

With dispatch_batches=False, an IterableDataset is wrapped in an IterableDatasetShard and loaded by a DataLoaderShard. For this path, two values that gather_for_metrics depends on come out wrong:

• DataLoaderShard.total_batch_size (L643-L649) reads the batch size from the plain torch BatchSampler. That sampler only knows the per-process batch size and has no num_processes, so the property returns 4 instead of 8 for batch_size=4 on 2 processes.
• DataLoaderStateMixin.begin() (L407) computes the length as getattr(self.dataset, "total_dataset_length", len(self.dataset)). IterableDatasetShard has no total_dataset_length, so it falls back to len(shard). That is the padded per-process length (8 for 12 samples), not the dataset length.

As a result remainder = 8 % 4 = 0, and gather_for_metrics does not drop the samples that were duplicated to fill the last batch.

# torchrun --nproc_per_node 2 repro.py   (CPU is enough)
import torch
from torch.utils.data import DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.utils import DataLoaderConfiguration


class SizedIterable(IterableDataset):
    def __init__(self, n):
        self.n = n

    def __iter__(self):
        yield from (torch.tensor(i) for i in range(self.n))

    def __len__(self):
        return self.n


for split_batches in (False, True):
    accelerator = Accelerator(
        cpu=True, dataloader_config=DataLoaderConfiguration(dispatch_batches=False, split_batches=split_batches)
    )
    for name, dataset in (("map-style", list(range(10))), ("iterable", SizedIterable(10))):
        dataloader = accelerator.prepare(DataLoader(dataset, batch_size=4))
        gathered = []
        for batch in dataloader:
            gathered += accelerator.gather_for_metrics(batch).tolist()
        accelerator.print(
            f"split_batches={split_batches} {name:9}: total_batch_size={dataloader.total_batch_size} "
            f"remainder={dataloader.remainder} gathered {len(gathered)} samples: {gathered}"
        )
    accelerator.state._reset_state()

Output on main:

split_batches=False map-style: total_batch_size=8 remainder=2 gathered 10 samples: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
split_batches=False iterable : total_batch_size=4 remainder=0 gathered 16 samples: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5]
split_batches=True map-style: total_batch_size=4 remainder=2 gathered 10 samples: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
split_batches=True iterable : total_batch_size=2 remainder=0 gathered 12 samples: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1]

For the same 10 samples, the sharded iterable loader returns 16 (or 12) samples from gather_for_metrics. Any metric computed from them counts the first samples twice.

The same SizedIterable(10) under the default dispatch_batches (the DataLoaderDispatcher path) returns [0..9], so dropping the padding for a sized iterable dataset is already supported there. Only the DataLoaderShard path gets it wrong. This is different from #3013, which was about streaming datasets without a known length. Here __len__ is defined.

Expected behavior

For a sized IterableDataset with dispatch_batches=False, the sharded dataloader should report the same values as the map-style path:

• total_batch_size == batch_size * num_processes (or batch_size with split_batches=True). This is what the DataLoaderShard docstring describes, and the same formula DataLoaderDispatcher.total_batch_size uses.
• total_dataset_length == len(dataset), and therefore remainder == len(dataset) % total_batch_size.

gather_for_metrics should then return exactly len(dataset) samples, as it does for the map-style and dispatched paths.

Suggested fix (small, in src/accelerate/data_loader.py):

1. In DataLoaderShard.total_batch_size, when self.dataset is an IterableDatasetShard, return dataset.batch_size if dataset.split_batches else dataset.batch_size * dataset.num_processes, the same formula as DataLoaderDispatcher.total_batch_size.
2. Give IterableDatasetShard a total_dataset_length property that returns len(self.dataset) (the unpadded inner length). begin() already looks for this attribute, and DataLoaderShard.total_dataset_length can use it too. Unsized datasets still raise inside begin()'s suppress, so they keep remainder == -1.

I have this working locally, with a unit test in tests/test_data_loader.py that compares the sharded iterable loader against the map-style loader for split_batches in {False, True}. If the fix approach looks right to you, could you assign this to me? I'll then open the PR.

AI assistance: drafted with an AI coding assistant (Claude). The reproduction above was run locally against the current default branch.