Skip to content

[python][torch] Support distributed iterable dataset sharding - #9429

Draft
XiaoHongbo-Hope wants to merge 11 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/torch-ddp-sharding
Draft

[python][torch] Support distributed iterable dataset sharding#9429
XiaoHongbo-Hope wants to merge 11 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/torch-ddp-sharding

Conversation

@XiaoHongbo-Hope

@XiaoHongbo-Hope XiaoHongbo-Hope commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

What changed

  • Shard streaming Torch datasets across DDP ranks and DataLoader workers when auto_detect_rank=True.
  • Keep rank sharding disabled by default because to_torch accepts caller-planned splits.
  • Resolve rank from torch.distributed, then torchrun environment variables, and preserve it in spawned workers.
  • Reject binding limits with multiple ranks.
  • Keep shuffled reads reproducible per epoch, rank, and worker.
  • Document uneven inputs and DistributedDataParallel.join().

Reference

The API is inspired by LanceDataset. Unlike Lance, this low-level API accepts external splits, so automatic sharding is opt-in.

Validation

  • 43 Torch tests passed.
  • CPU torchrun test executes DDP forward/backward with 2 ranks and 2 spawned DataLoader workers.
  • Split assignment is complete and non-overlapping.
  • Pre-sharded splits remain unchanged by default.
  • flake8 and git diff --check passed.

@XiaoHongbo-Hope
XiaoHongbo-Hope force-pushed the codex/torch-ddp-sharding branch from fb7c936 to 7dcf4cf Compare August 27, 2026 11:25
@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as ready for review August 29, 2026 08:51
@XiaoHongbo-Hope
XiaoHongbo-Hope force-pushed the codex/torch-ddp-sharding branch from 49b5d04 to 3334cd9 Compare August 29, 2026 09:16
@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as draft August 29, 2026 10:05
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