Skip to content

Fix duplicate samples in multi-worker streaming dataloaders - #203

Open
Travor278 wants to merge 3 commits into
huggingface:mainfrom
Travor278:fix-issue-200-worker-sharding
Open

Travor278 wants to merge 3 commits into
huggingface:mainfrom
Travor278:fix-issue-200-worker-sharding

Conversation

@Travor278

Copy link
Copy Markdown

This PR fixes the worker-sharding bug from #200.

When VQADataset.iter_for_worker() was used with a streaming dataset and DataLoader(num_workers > 1), each worker would walk the full stream independently. That meant every sample was processed once per worker instead of once overall.

What changed:

  • shard streaming datasets per worker via .shard(...) when the underlying dataset supports it
  • fall back to stride-based sharding for iterable datasets that do not implement .shard()
  • add a real multi-worker regression test that reproduces the old duplication and verifies the fixed behavior
  • add a small benchmark and a training-path smoke test to compare duplicate counts and effective throughput before vs after the fix

Why this matters:

  • it removes duplicated samples during multi-worker streaming training
  • it makes effective throughput line up with actual useful sample coverage
  • it keeps the fix robust even when the iterable dataset is not a Hugging Face streaming object

Validation:

  • D:\Dev\conda-envs\py313\python.exe -m unittest tests.test_worker_sharding -v
  • D:\Dev\conda-envs\py313\python.exe tests\benchmark_worker_sharding.py --workers 4 --samples 2000 --repeats 5
    • before: 40000 processed / 10000 unique / 30000 duplicates / 554.33 effective samples/s
    • after: 10000 processed / 10000 unique / 0 duplicates / 900.49 effective samples/s
  • D:\Dev\conda-envs\py313\python.exe tests\smoke_worker_sharding_train_path.py --device auto
    • legacy path: 288 total / 96 unique / 192 duplicates / 3.00x
    • fixed path: 96 total / 96 unique / 0 duplicates / 1.00x

Fixes #200.

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.

VQADataset.iter_for_worker() lacks explicit worker sharding, causing data duplication in multi-worker DataLoader

1 participant