Repository navigation
Stop BatchSamplerShard padding from mutating sampler-owned batches - #4507
Open
inchang-ing wants to merge 1 commit into
Open
inchang-ing wants to merge 1 commit into
inchang-ing wants to merge 1 commit into
Conversation
_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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #4505
Problem
BatchSamplerShardpads lists by mutating them in place, and those lists can be sampler-owned:_iter_with_splitbindsinitial_datato the sampler's first batch list (initial_data = batch, line 187) and then doubles it in place (initial_data += initial_datain the degenerate-case padding);_iter_with_no_splitcan reachbatch += initial_data[cycle_index:end_index]whilebatchis 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
[[0, 1, 0, 1]]after one pass; post-fix it stays[[0, 1]]and the original loader still yields[[0, 1]].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 --checkclean on both touched files;ruff checkfinding counts unchanged vsmainfordata_loader.py.Not run / limitations
split_batches=True, multi-rank joins) are exercised by the existingtest_data_loader.pysuite, 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.