Deterministic distributed sampler that streams across epoch boundaries. Each epoch uses the first ``total_size`` samples from a deterministic shuffle of the dataset. The sampler can produce a fixed number of per-rank samples (``num_samples``) starting from an arbitrary per-rank off
| 60 | yield |
| 61 | |
| 62 | class StatelessResumableDistributedSampler(Sampler): |
| 63 | """Deterministic distributed sampler that streams across epoch boundaries. |
| 64 | |
| 65 | Each epoch uses the first ``total_size`` samples from a deterministic |
| 66 | shuffle of the dataset. The sampler can produce a fixed number of |
| 67 | per-rank samples (``num_samples``) starting from an arbitrary per-rank |
| 68 | offset, transparently crossing epoch boundaries with fresh shuffles. |
| 69 | |
| 70 | When ``num_samples`` is *None* (default) the sampler yields the remaining |
| 71 | samples in the current epoch — this preserves backward compatibility with |
| 72 | code that rebuilds the dataloader at every epoch boundary. |
| 73 | """ |
| 74 | |
| 75 | def __init__( |
| 76 | self, |
| 77 | dataset, |
| 78 | num_replicas: int, |
| 79 | rank: int, |
| 80 | total_size: int, |
| 81 | seed: int = 42, |
| 82 | start_global_offset_samples: int = 0, |
| 83 | num_samples: int | None = None, |
| 84 | ): |
| 85 | assert start_global_offset_samples >= 0, "start_global_offset_samples must be >= 0" |
| 86 | self.dataset = dataset |
| 87 | self.num_replicas = num_replicas |
| 88 | self.rank = rank |
| 89 | self.total_size = int(total_size) |
| 90 | self.seed = int(seed) |
| 91 | self.dataset_size = len(self.dataset) |
| 92 | assert self.dataset_size > 0, "dataset must have positive length" |
| 93 | assert self.total_size > 0, "total_size must be > 0" |
| 94 | assert self.total_size <= self.dataset_size, ( |
| 95 | f"total_size ({self.total_size}) cannot exceed dataset size ({self.dataset_size})" |
| 96 | ) |
| 97 | assert self.total_size % self.num_replicas == 0, ( |
| 98 | f"total_size ({self.total_size}) must be divisible by num_replicas ({self.num_replicas})" |
| 99 | ) |
| 100 | assert num_samples is None or num_samples >= 0, "num_samples must be >= 0" |
| 101 | |
| 102 | self.per_rank_len_per_epoch = self.total_size // self.num_replicas |
| 103 | self._global_offset = int(start_global_offset_samples) |
| 104 | self._num_samples = num_samples |
| 105 | |
| 106 | def __len__(self): |
| 107 | if self._num_samples is not None: |
| 108 | return self._num_samples |
| 109 | mod = self._global_offset % self.per_rank_len_per_epoch |
| 110 | return self.per_rank_len_per_epoch - mod if mod != 0 else self.per_rank_len_per_epoch |
| 111 | |
| 112 | def _epoch_perm(self, epoch_idx: int): |
| 113 | g = torch.Generator() |
| 114 | g.manual_seed(self.seed + epoch_idx) |
| 115 | return torch.randperm(self.dataset_size, generator=g).tolist()[: self.total_size] |
| 116 | |
| 117 | def _epoch_slice_for_rank(self, perm): |
| 118 | return perm[self.rank : self.total_size : self.num_replicas] |
| 119 |
no outgoing calls
no test coverage detected