fix: shard TorchDataNodeDataLoader via DistributedSampler under DDP#106
Open
felixmaximilian wants to merge 2 commits into
Open
fix: shard TorchDataNodeDataLoader via DistributedSampler under DDP#106felixmaximilian wants to merge 2 commits into
felixmaximilian wants to merge 2 commits into
Conversation
The loader hard-codes RandomSampler and is not a torch DataLoader, so pytorch-lightning cannot inject a DistributedSampler. Under devices>1 every rank draws an identical (same-seed) sample sequence: the all-reduced gradient degenerates to a single rank's — N times the compute for an unchanged effective batch, silently. - build the sampler/node graph lazily at first iter()/len(): framework integrations construct the loader before the process group exists, so the sharding decision must be taken at iteration time (with a rebuild guard if the group appears after a premature build) - use DistributedSampler (disjoint shards; world_size x per-rank batch = effective batch) when torch.distributed is initialized; exact previous behavior otherwise - expose set_epoch() for per-epoch reshuffling (callers: see rmind's DistributedSamplerEpochSetter callback) - reject method='process' under distributed training: workers fork after CUDA init and deadlock Verified in rmind 2-GPU training on 2x RTX 4090 (both ranks log disjoint shard_len = N/2 and train to completion) plus gloo-backend unit tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The DDP-sharding change carried pyright ignore comments, but CI's 'just check' type-checks with ty, which reads only '# ty: ignore[rule]'. Add ty pragmas at the torchdata/torch stub limitations (ParallelMapper kwargs unpack, sampler dataset protocols, prefetch_factor arithmetic, the deliberate _loader poke in the DDP rebuild test) and at the pre-existing HydraConfig(target=...) alias sites ty now flags; drop four ignore comments the call-site fix made unused. No runtime changes; ruff/format/ty all green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Author
|
@naidenovaleksei This answers your question why two GPUs were as fast as one wrt iter/s! |
Contributor
|
iirc |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
TorchDataNodeDataLoaderhard-codesRandomSamplerand is not atorch.utils.data.DataLoader, so pytorch-lightning cannot inject aDistributedSampler. Undertrainer.devices>1, Lightning seeds all ranks identically → every rank draws the same sample sequence. The all-reduced gradient degenerates to a single rank's: N× the compute for an unchanged effective batch — silently.We hit this in rmind DDP training (yaak-ai/rmind
feat/throttle-rebalance-otf): first 2-GPU smoke showed both ranks iterating the identical unsharded 549-batch epoch. An audit of the W&B project found no priordevices>1runs on this loader, so no historical results are affected — the trap was latent.Fix
iter()/len()— framework integrations (hydra + Lightning) instantiate the loader beforetrainer.fitinitializes the process group, so an init-timeis_initialized()check can never shard rank 0. A rebuild guard re-shards if the process group appears after a premature build.DistributedSampler(disjoint shards;world_size × per-rank batch = effective batch) whenevertorch.distributedis initialized; bit-identical previous behavior otherwise.set_epoch()exposed for per-epoch reshuffling (Lightning only forwards epochs to plain DataLoaders; rmind ships a small callback calling this).method='process'rejected under distributed training: ParallelMapper workers fork after CUDA init and deadlock (known failure, previously observed single-GPU with the shared-memory TensorDict).Verification
distributed sampler rank={0,1} world_size=2 shard_len=N/2and train normally.set_epochwith one, late-group rebuild guard, process-method rejection. (test_dataloaders[{thread,process}-yaak_dataset]fail identically on cleanmainin my environment — pre-existing, unrelated.)🤖 Generated with Claude Code