feat: DDP-sharded dataloader + trainer=multi_gpu (companion to rbyte#106)

[felixmaximilian]

Companion to yaak-ai/rbyte#106 — the two should land together.

Problem

rbyte's TorchDataNodeDataLoader hard-codes RandomSampler and is not a torch.utils.data.DataLoader, so Lightning cannot inject a DistributedSampler. Under trainer.devices>1 every rank draws an identical (same-seed) batch sequence: the all-reduced gradient degenerates to a single rank's, and an epoch takes the same number of steps as on one GPU. This is the long-observed "2 GPUs = 1 GPU it/s" — N× the compute for nothing. (W&B audit: no completed devices>1 runs exist, so no results are affected; the trap was latent.)

Changes

• rmind.components.dataloader.DistributedTorchDataNodeDataLoader — drop-in for the rbyte loader; shards via DistributedSampler when torch.distributed is initialized (world_size × per-rank batch = effective batch), exact rbyte behavior otherwise. Built lazily at first iter()/len(): hydra constructs loaders before trainer.fit initializes the process group, so an init-time check can never shard rank 0 (verified empirically). Rebuild guard if the group appears after a premature build; method='process' rejected under DDP (fork-after-CUDA deadlock).
• DistributedSamplerEpochSetter callback (per-epoch reshuffling; no-op on single GPU) registered in the pretrain/finetune callback lists.
• trainer=multi_gpu config group: devices: 2, strategy: ddp_find_unused_parameters_true (frozen-path modules leave grad-less params each step). Halve per-rank batch_size to keep the effective batch — 2×32 ≡ bs64.
• yaak datamodules point at the new loader.

Interim plan

Once rbyte#106 is released and the pin bumps, drop the loader class here and revert the datamodule targets; the epoch-setter callback (Lightning-specific) and trainer=multi_gpu stay.

Verification

2-GPU run on 2× RTX 4090 (run ylpsmyv9): both ranks log distributed sampler rank={0,1} world_size=2 with disjoint shard_len = N/2, effective-batch-64 math preserved, ~2× epoch wall-clock vs single GPU. Unit tests for the sharding logic live in rbyte#106 (gloo backend).

🤖 Generated with Claude Code