JingsongLi commented on code in PR #9429: URL: https://github.com/apache/paimon/pull/9429#discussion_r3888762800
########## docs/docs/pypaimon/pytorch.md: ########## @@ -60,6 +60,30 @@ when it is false, it will read the full amount of data into memory. **`prefetch_concurrency`** (default: 1): In streaming row mode, controls reader threads per DataLoader worker. It has no effect in non-streaming mode. +### Distributed Sharding + +Streaming reads shard splits across DDP ranks and DataLoader workers: + +```python +dataset = table_read.to_torch( + splits, + streaming=True, + auto_detect_rank=True, +) +dataloader = DataLoader(dataset, batch_size=32, num_workers=2) Review Comment: **[P1] Use a spawn-safe DataLoader in this DDP example** On Linux, leaving `multiprocessing_context` unset makes a multi-worker DataLoader use `fork`. PyTorch explicitly warns that DDP with NCCL (and Gloo over InfiniBand) is not fork-safe and is likely to deadlock. This is especially easy to miss because the new torchrun test uses `multiprocessing_context="spawn"`, so it does not exercise the documented path. Please set `multiprocessing_context="spawn"` (or `forkserver`) in the example and mention the usual guarded-entry-point requirement for spawn. See https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html. ########## paimon-python/pypaimon/read/datasource/torch_dataset.py: ########## @@ -40,6 +41,59 @@ def _share_epoch_with_torch_workers(value): return torch.tensor(value, dtype=torch.long).share_memory_() +def _validate_distributed_context(rank: int, world_size: int): + if isinstance(rank, bool) or not isinstance(rank, int): + raise ValueError("rank must be an int") + if isinstance(world_size, bool) or not isinstance(world_size, int): + raise ValueError("world_size must be an int") + if world_size <= 0: + raise ValueError("world_size must be greater than 0") + if rank < 0 or rank >= world_size: + raise ValueError("rank must satisfy 0 <= rank < world_size") + return rank, world_size + + +def _resolve_distributed_context(auto_detect_rank: bool): + if not isinstance(auto_detect_rank, bool): + raise ValueError("auto_detect_rank must be a bool") + if not auto_detect_rank: + return 0, 1 + + distributed = getattr(torch, "distributed", None) + if ( + distributed is not None + and distributed.is_available() + and distributed.is_initialized() + ): + rank = distributed.get_rank() Review Comment: **[P2] Shard against the DDP process group, not always the global world** Calling `get_rank()` and `get_world_size()` without a group always resolves the default global process group, while DDP may be constructed with `process_group=subgroup`. In subgroup or hybrid-parallel training, `auto_detect_rank=True` will therefore shard inputs across the global world instead of the replicas synchronized by DDP; a subgroup can see only part of its dataset, or model-parallel ranks can receive different inputs. Please either accept explicit sharding rank/world size resolved from the intended DDP group before serializing the dataset to workers, or document that auto-detection supports only whole-world DDP and require caller-side pre-sharding otherwise. See https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
