| 44 | |
| 45 | # Subclass DistributedSampler to use PCG64DXSM for shuffling |
| 46 | class DistributedSamplerPCG64DXSM(DistributedSampler): |
| 47 | def __iter__(self) -> Iterator[int]: |
| 48 | if self.shuffle: |
| 49 | # deterministically shuffle based on epoch and seed |
| 50 | # use numpy's RNG PCG64DXSM instead of torch.randperm |
| 51 | rng = np.random.Generator(np.random.PCG64DXSM(self.seed + self.epoch)) |
| 52 | indices = rng.permutation(len(self.dataset)).tolist() # type: ignore[arg-type] |
| 53 | else: |
| 54 | indices = list(range(len(self.dataset))) # type: ignore[arg-type] |
| 55 | |
| 56 | if not self.drop_last: |
| 57 | # add extra samples to make it evenly divisible |
| 58 | padding_size = self.total_size - len(indices) |
| 59 | if padding_size <= len(indices): |
| 60 | indices += indices[:padding_size] |
| 61 | else: |
| 62 | indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size] |
| 63 | else: |
| 64 | # remove tail of data to make it evenly divisible. |
| 65 | indices = indices[: self.total_size] |
| 66 | assert len(indices) == self.total_size |
| 67 | |
| 68 | # subsample |
| 69 | indices = indices[self.rank : self.total_size : self.num_replicas] |
| 70 | assert len(indices) == self.num_samples |
| 71 | |
| 72 | return iter(indices) |
| 73 | |
| 74 | |
| 75 | def build_tokenizer( |
no outgoing calls
no test coverage detected