Skip to content

Avoid mutating sampler-owned batches during distributed padding - #4506

Open
davrens101 wants to merge 1 commit into
huggingface:mainfrom
davrens101:retain-sampler-batches-round5
Open

davrens101 wants to merge 1 commit into
huggingface:mainfrom
davrens101:retain-sampler-batches-round5

Conversation

@davrens101

@davrens101 davrens101 commented Nov 6, 2026 •

Copy link
Copy Markdown

What does this PR do?

Fixes #4505.

Keep padding local to BatchSamplerShard rather than extending lists retained by the underlying sampler. Currently, preparing and iterating a loader with precomputed batches can change the original loader's future samples: [0, 1] becomes [0, 1, 0, 1]. The unsplit path similarly changes a seven-sample original stream into eight samples.

Replace two in-place concatenations in the padding paths with non-mutating concatenations. Padded shard contents and ordering stay the same; full batches do not acquire unconditional copies. The docstring now records the sampler-ownership invariant.

This is separate from #4395's reported shard-length correction: no length calculation changes here.

Verification

  • Both new behavioral regression cases fail on untouched main f47439c1fb24721ea71ab98f9b512726b5a97828, and pass with this fix. They check actual prepared/original loader sample streams, both ranks and repeated iterations, for split and whole-batch padding.
  • Entire tests/test_data_loader.py: 42 passed, including existing stateful-loader cases.
  • Repository-wide Ruff 0.13.1 lint and formatting: pass, 183 files formatted.
  • Actual two-process CPU/Gloo DDP with Accelerator.prepare, backward and an SGD step: both padding modes match an independent analytic global-batch gradient on both ranks. Weights are 0.099999994 vs expected 0.1 (split) and 0.704999983 vs expected 0.705 (whole-batch). The original loaders still yield exactly 2 and 7 samples. On untouched main, all four corresponding rank/mode runs fail the original-sample invariant.

The Windows torchrun elastic rendezvous initially hit this build's missing libuv support. The successful DDP smoke launches two real processes directly and initializes Gloo with an explicit non-libuv TCP rendezvous. No GPU execution is claimed.

Before submitting

AI-assisted investigation, implementation and independent agent review. No claim of human review of this specific patch.

Upstream CI status

The posted head is fabb3f3d13f3df391894f399f90237828bf1a8a4. Run Tests 42300201084, Quality Check 42300201286, Run Import Tests 42300201235, and Build PR Documentation 42300201998 report action_required, with zero jobs started. Upstream tests are not claimed passing; the local affected suite and both actual two-process training modes passed again after publication.

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