Skip to content

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

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

[python][torch] Support distributed iterable dataset sharding#9429
XiaoHongbo-Hope wants to merge 7 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.
  • Resolve the DDP context at iteration time from torch.distributed, then torchrun environment variables.
  • Reject a binding read limit with multiple ranks to avoid uneven DDP input.
  • Make shuffled reads reproducible per epoch, rank, and worker.
  • Allow rank detection to be disabled with auto_detect_rank=False.

Reference

The auto_detect_rank API and its default follow Lance Torch LanceDataset.

Validation

  • 41 Torch tests passed.
  • CPU torchrun test passed with 2 ranks and 2 DataLoader workers, with complete and non-overlapping split assignment.
  • 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
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