Skip to content

Handle unknown batch size in sharded dataloader metadata - #4491

Draft
Aurantium-47 wants to merge 1 commit into
huggingface:mainfrom
Aurantium-47:fix/nullable-variable-batch-metadata
Draft

Aurantium-47 wants to merge 1 commit into
huggingface:mainfrom
Aurantium-47:fix/nullable-variable-batch-metadata

Conversation

@Aurantium-47

Copy link
Copy Markdown

What does this PR do?

Return None for DataLoaderShard.total_batch_size when a variable batch sampler has batch_size=None, instead of raising TypeError while multiplying it by the process count. This applies after existing SkipBatchSampler unwrapping.

Keep the remainder explicitly unknown (-1) when no constant batch size is available, document that contract, and correct the gather_for_metrics informational message: an unknown remainder does not necessarily mean that the dataset has no length. Gathering/trimming logic is unchanged; callers remain responsible for duplicate removal when the remainder is unknown.

Known-size behavior is unchanged. This does not add variable split-batch support, change sharding/padding, or make uneven distributed collectives safe.

Validation

Reapplied the independently reviewed patch to upstream main f47439c1fb24721ea71ab98f9b512726b5a97828 and reran:

  • PYTHONPATH=src ACCELERATE_USE_CPU=true python -m pytest tests/test_variable_batch_metadata.py tests/test_data_loader.py -q -rs: 41 passed, 12 skipped (optional torchdata.stateful_dataloader absent).
  • make quality with the project's pinned Ruff 0.13.1: all checks passed, 184 files formatted.
  • Additional local probe: 12 odd-tail, two simulated-rank, drop-last and nested-skip combinations passed; iterable smoke checks passed.
  • The 13 new tests had 8 expected failures and 5 passes against the original unchanged source before this fix. Independent automated review reran the focused suite and reviewed the production diff.

Linux CPU, Python 3.12.14, PyTorch 2.15.0+cpu. No distributed process group was launched for this validation; DDP, GPU, XLA and full repository-suite validation are not claimed.

Related but distinct: #4416 repairs sized iterable metadata; #4489 chooses sampler length for the remainder. Neither fixes batch_size=None. A local compatibility probe with their fetched changes passed (45 passed, 12 optional skips), preserving #4489's unsized-sampler fallback. Those changes are not included here; its overlapping modulo line will need care if it lands first.

Developed with AI assistance, including separate automated code review. No human review or prior maintainer approval is claimed.

Before submitting

  • Read the contributor guideline and Pull Request section.
  • Discussed/approved through an issue or forum (not claimed; submitted as a focused draft proposal).
  • Updated the affected metadata/remainder docstrings.
  • Added necessary regression tests.

Who can review?

Submitted as a draft; no reviewer is assigned by this description.

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