Fix gather_for_metrics remainder when a sized sampler selects a subset

[Eryx-Z]

Summary

Fixes #4410.

DataLoaderStateMixin.begin() computed the final-batch remainder from len(dataset). When a sized sampler (e.g. SubsetRandomSampler) selects fewer indices than the underlying map-style dataset, gather_for_metrics then drops or duplicates samples.

Change: prefer len(batch_sampler.sampler) (then nested / top-level sampler) when sized; otherwise keep total_dataset_length / len(dataset). After Accelerator.prepare, the real selection often lives on batch_sampler.sampler while self.sampler can be a default SequentialSampler over the full dataset — so that order matters.

Validation

pytest tests/test_data_loader.py -k "remainder_uses_sampler or skip_first_batches_preserves" -q
# 5 passed

New test reproduces the issue's SubsetRandomSampler case and the Subset control.

Coordination

Comment left on #4410: issue was unassigned with no open PR; happy to close if @gorkadabo1 prefers to take it.

[Michael-RDev]

With BatchSampler(drop_last=True), DataLoader.drop_last stays False. Selecting 10 samples with batch size 4 correctly produces 8 samples but this change makes gather_for_metrics keep only 6. Could we check the batch sampler’s setting and add a test?

[Eryx-Z]

Good catch — fixed in 49c2569.

begin() now uses _effective_drop_last() which also walks batch_sampler.drop_last (and nested batch_sampler.batch_sampler), so when BatchSampler(drop_last=True) is passed and DataLoader.drop_last stays False, remainder stays -1 and gather_for_metrics keeps all 8 yielded samples.

Added test_remainder_respects_batch_sampler_drop_last covering the 10→8 case.

[Michael-RDev]

begin() runs before set_epoch(), so this can use the previous epoch’s sampler length. I reproduced a sampler growing from 5 to 6 samples where the sixth sample gets dropped from metrics. Could we also set the epoch first and cover this in a test?

[Eryx-Z]

Agreed — fixed in the same commit.

DataLoaderShard.__iter__ and DataLoaderDispatcher.__iter__ now call set_epoch(self.iteration) before begin(), so remainder uses the current epoch’s sampler length.

Added test_remainder_uses_length_after_set_epoch with a growing sampler (len 5 → 6 across iterations); second pass now keeps all 6 samples.

[Michael-RDev]

This fixes the reported subset-sampler case but I found two cases where valid samples get dropped from metrics: batch samplers using ``drop_last=True` and samplers whose length changes between epochs.

Both reproductions pass on the parent commit and fail here. I’ve left some details, please address them and add regression tests :)

[Eryx-Z]

Addressed both review notes from @Michael-RDev in 49c2569:

1. BatchSampler.drop_last=True — _effective_drop_last() honors the batch sampler’s drop_last so remainder is not computed when incomplete batches are dropped.
2. set_epoch before begin — both shard and dispatcher iterators now set the epoch first, then compute remainder.

Local: pytest tests/test_data_loader.py -k "remainder_uses_sampler or remainder_respects_batch or remainder_uses_length_after or skip_first_batches_preserves" → 7 passed.

[Eryx-Z]

Friendly bump for maintainers: both review notes from @Michael-RDev are addressed in 49c2569 (BatchSampler.drop_last via _effective_drop_last(), and set_epoch before begin), with regression tests included. Ready for another look whenever you have a moment.