dataloader_rs/sampler/
distributed.rs

1use crate::sampler::Sampler;
2
3/// Shards indices across `world_size` workers for distributed training.
4pub struct DistributedSampler<S: Sampler> {
5    inner: S,
6    rank: usize,
7    world_size: usize,
8}
9
10impl<S: Sampler> DistributedSampler<S> {
11    /// Create a new `DistributedSampler`.
12    ///
13    /// # Panics
14    ///
15    /// Panics if `rank >= world_size` or `world_size == 0`.
16    pub fn new(inner: S, rank: usize, world_size: usize) -> Self {
17        assert!(world_size > 0, "world_size must be > 0");
18        assert!(rank < world_size, "rank must be < world_size");
19        Self {
20            inner,
21            rank,
22            world_size,
23        }
24    }
25}
26
27impl<S: Sampler> Sampler for DistributedSampler<S> {
28    fn indices(&mut self, dataset_len: usize) -> Vec<usize> {
29        let all = self.inner.indices(dataset_len);
30        // As in PyTorch: pad the epoch to a multiple of `world_size` by
31        // wrapping around, so every rank gets as many indices, and give this
32        // rank every `world_size`-th position from `rank`.
33        (0..all.len().div_ceil(self.world_size))
34            .map(|i| all[(self.rank + i * self.world_size) % all.len()])
35            .collect()
36    }
37
38    fn len(&self, dataset_len: usize) -> usize {
39        self.inner.len(dataset_len).div_ceil(self.world_size)
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use crate::sampler::{DistributedSampler, RandomSampler, Sampler, SequentialSampler};
46
47    #[test]
48    fn distributed_partitions_have_expected_total_len() {
49        let world_size = 3;
50        let n = 10;
51        let mut all_indices: Vec<usize> = Vec::new();
52        for rank in 0..world_size {
53            let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
54            all_indices.extend(ds.indices(n));
55        }
56        assert_eq!(all_indices.len(), 12);
57    }
58
59    #[test]
60    fn distributed_equal_length_per_rank() {
61        let world_size = 4;
62        let n = 10;
63        let lengths: Vec<usize> = (0..world_size)
64            .map(|rank| {
65                let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
66                ds.indices(n).len()
67            })
68            .collect();
69        assert!(lengths.windows(2).all(|w| w[0] == w[1]));
70        assert_eq!(lengths[0], 3);
71    }
72
73    #[test]
74    fn distributed_len_matches_indices() {
75        for (n, world_size) in [(10, 3), (10, 4), (8, 4), (1, 4), (0, 2)] {
76            for rank in 0..world_size {
77                let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
78                assert_eq!(ds.len(n), ds.indices(n).len(), "n={n} world={world_size}");
79            }
80        }
81    }
82
83    #[test]
84    fn distributed_pads_by_wrapping_around() {
85        let shard = |n, rank, world_size| {
86            DistributedSampler::new(SequentialSampler, rank, world_size).indices(n)
87        };
88        // 10 indices padded to 12 with 0 and 1.
89        assert_eq!(shard(10, 0, 4), vec![0, 4, 8]);
90        assert_eq!(shard(10, 2, 4), vec![2, 6, 0]);
91        assert_eq!(shard(10, 3, 4), vec![3, 7, 1]);
92        // Padding longer than the epoch wraps more than once.
93        for rank in 0..4 {
94            assert_eq!(shard(1, rank, 4), vec![0]);
95        }
96        assert_eq!(shard(0, 1, 4), Vec::<usize>::new());
97    }
98
99    #[test]
100    fn distributed_exact_divisor() {
101        let world_size = 4;
102        let n = 8;
103        for rank in 0..world_size {
104            let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
105            assert_eq!(ds.indices(n).len(), 2);
106        }
107    }
108
109    #[test]
110    #[should_panic]
111    fn distributed_rank_out_of_bounds_panics() {
112        DistributedSampler::new(SequentialSampler, 4, 4);
113    }
114
115    #[test]
116    #[should_panic]
117    fn distributed_zero_world_size_panics() {
118        DistributedSampler::new(SequentialSampler, 0, 0);
119    }
120
121    #[test]
122    fn distributed_wraps_random_sampler() {
123        let world_size = 2;
124        let n = 6;
125        let mut ds = DistributedSampler::new(RandomSampler::new(99), 0, world_size);
126        let indices = ds.indices(n);
127        assert_eq!(indices.len(), n / world_size);
128        assert!(indices.iter().all(|&i| i < n));
129    }
130}
dataloader_rs/sampler/distributed.rs

dataloader_rs/sampler/
distributed.rs
