Skip to content

Stop BatchSamplerShard padding from mutating sampler-owned batches - #4507

Open
inchang-ing wants to merge 1 commit into
huggingface:mainfrom
inchang-ing:persist-sampler-batches-4505
Open

inchang-ing wants to merge 1 commit into
huggingface:mainfrom
inchang-ing:persist-sampler-batches-4505

Conversation

@inchang-ing

Copy link
Copy Markdown

Closes #4505

Problem

BatchSamplerShard pads lists by mutating them in place, and those lists can be sampler-owned:

  • _iter_with_split binds initial_data to the sampler's first batch list (initial_data = batch, line 187) and then doubles it in place (initial_data += initial_data in the degenerate-case padding);
  • _iter_with_no_split can reach batch += initial_data[cycle_index:end_index] while batch is still the sampler's retained final partial batch.

For samplers that retain their batches (e.g. precomputed batches), one pass over a prepared loader changes the original loader's future samples: the issue's two-sample loader becomes [0, 1, 0, 1] after a single split/padded pass, and the unsplit path extends the sampler-owned final partial batch in place.

Fix

Rebind instead of extending in place at the two mutation sites (initial_data = initial_data + initial_data, batch = batch + initial_data[cycle_index:end_index]). The prepared view is unchanged — same padded batches are yielded — but sampler-owned lists are never touched.

Verification

  • The issue's exact reproduction: pre-fix the sampler's retained batch becomes [[0, 1, 0, 1]] after one pass; post-fix it stays [[0, 1]] and the original loader still yields [[0, 1]].
  • Regression test tests/test_data_loader.py::test_batch_sampler_shard_does_not_mutate_sampler_batches (asserts both the retained batches and the original loader are unchanged after a prepared pass): fails on unpatched code (verified via stash A/B), full file 28 passed / 13 skipped (GPU-only skips) with the fix.
  • ruff format --check clean on both touched files; ruff check finding counts unchanged vs main for data_loader.py.

Not run / limitations

  • Verified CPU/single-process; distributed combination branches (split_batches=True, multi-rank joins) are exercised by the existing test_data_loader.py suite, which passes unchanged.

Prepared with the assistance of an AI coding agent (an automated pipeline working on behalf of the account owner); the account owner submits and stands behind this PR.

_iter_with_split bound initial_data to the sampler's first batch list
and then doubled it in place, and _iter_with_no_split could extend the
sampler-owned final partial batch via batch += ... — both changed what
the original sampler (and through a retained-batch sampler, the original
loader) would yield afterwards.

Rebind the padded lists instead of extending in place; the prepared view
is unchanged.

Closes huggingface#4505
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.

BatchSamplerShard padding mutates retained sampler batches and changes original loader samples

1 participant