Repository navigation
Fix split-batch shard length for uneven tails - #4395
Open
AaravMehta-07 wants to merge 1 commit into
Open
AaravMehta-07 wants to merge 1 commit into
AaravMehta-07 wants to merge 1 commit into
Conversation
Assisted-by: OpenAI Codex
Author
|
Hi maintainers, whenever you have a chance, I'd appreciate a review of this fix for matching split-batch shard lengths to uneven rank-local tails. The focused batch-sampler cases passed locally (11 tests), including the new length regressions and custom-sampler compatibility. Please let me know if you'd prefer a different approach or additional focused coverage. Thank you for your time. AI-assisted follow-up. |
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.
Summary
Correct
BatchSamplerShard.__len__whensplit_batches=True,even_batches=False, and a final partial batch does not reach a rank's slice. Iteration already skips that rank's empty tail, but the reported length still includes it. For two samples, batch size four and two processes, the reported shard lengths are[1, 1]although iteration yields[1, 0]. This can give callers an incorrect dataloader length.Derive the final partial batch size from the underlying sampler for the exact ordinary PyTorch
BatchSampler, and remove the unavailable tail from that rank's length. No iteration or data ordering changes. Custom subclasses retain their existing length behavior because their batching or sampler sizing may differ; a control covers a subclass with an unsized iterator.Validation
ccfdff7: fails with[1, 1] != [1, 0].tests/test_data_loader.py: 30 passed, 12 skipped in the CPU Windows environment.git diff --checkpasses.Tests use actual PyTorch samplers and library implementations, with no distributed-launch or GPU claim. The local Datasets 5.1.0/dill stack initially failed in an unrelated Arrow fingerprint test; an isolated supported dependency overlay (Datasets 4.3.0, dill 0.4.0, PyArrow 22.0.0) allowed the entire affected file to pass. No dependency pins change in this PR.
Duplicate checks found no matching active split-batch length fix. Existing #4389 concerns shared modules in regional compilation and is untouched.
AI disclosure: investigated and implemented with OpenAI Codex (GPT-6 Luna worker), with an independent root audit, a custom-subclass compatibility correction, baseline replay and full affected-file test execution.