Skip to content

Fix split-batch shard length for uneven tails - #4395

Open
AaravMehta-07 wants to merge 1 commit into
huggingface:mainfrom
AaravMehta-07:codex/accelerate-offline-fix
Open

AaravMehta-07 wants to merge 1 commit into
huggingface:mainfrom
AaravMehta-07:codex/accelerate-offline-fix

Conversation

@AaravMehta-07

Copy link
Copy Markdown

Summary

Correct BatchSamplerShard.__len__ when split_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

  • Regression run against upstream ccfdff7: fails with [1, 1] != [1, 0].
  • Independently rerun full tests/test_data_loader.py: 30 passed, 12 skipped in the CPU Windows environment.
  • Worker checked ordinary sampler combinations across sample counts 0-20, batch sizes 1-6, process counts 1-4, split/even modes and drop-last settings: no length/iteration mismatches after the fix.
  • Pinned changed-file pre-commit checks pass (Ruff lint/format and merge-conflict check; YAML has no matching changed files).
  • git diff --check passes.

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.

@AaravMehta-07

Copy link
Copy Markdown
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.

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