Skip to content

Fix metric padding removal after skip_first_batches - #4235

Open
Excelius-Wang wants to merge 1 commit into
huggingface:mainfrom
Excelius-Wang:fix-skip-batches-metric-padding
Open

Fix metric padding removal after skip_first_batches#4235
Excelius-Wang wants to merge 1 commit into
huggingface:mainfrom
Excelius-Wang:fix-skip-batches-metric-padding

Conversation

@Excelius-Wang

Copy link
Copy Markdown

What does this PR do?

skip_first_batches wraps the batch sampler in SkipBatchSampler, which does not expose batch_size. This makes DataLoaderShard.total_batch_size raise; the suppressed exception leaves remainder=-1, so gather_for_metrics keeps duplicated padding samples in the final batch.

With 10 samples, two processes and a per-process batch size of 2, normal iteration gathers 10 samples. Calling skip_first_batches(loader, 0) instead gathers 12, including duplicated samples 0 and 1.

Unwrap skip samplers before reading the existing batch-size metadata, and preserve _drop_last when reconstructing the loader. This also handles nested skipping without changing sampling order.

Validation

  • 66 related tests passed, 16 skipped.
  • All 18 new parameterized regressions fail on the baseline and pass with the fix.
  • Actual CPU/Gloo testing: the expanded 48-case two-process matrix passed twice; the earlier 24-case three-process matrix also passed twice. Coverage includes split batches, different batch sizes, divisible and non-divisible dataset sizes, drop_last, and zero/one/nested skips.
  • The distributed regression fails on the baseline.
  • Full make quality and git diff --check pass.

GPU/TPU execution was not tested. Changes were prepared with AI assistance and reviewed by the contributor.

Before submitting

  • Read the contribution guidelines.
  • Added regression tests.
  • Prior issue discussion: none.
  • Documentation changes: not needed; this restores existing behavior.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant