| 244 | |
| 245 | |
| 246 | class FinetuneDistSampler(Sampler): |
| 247 | # Distrubuted Sampler ensuring data in a batch are of the same type (e.g. text, image-text) |
| 248 | def __init__(self, dataset: FinetuneDataset, num_replicas: Optional[int] = None, |
| 249 | rank: Optional[int] = None, shuffle: bool = True, |
| 250 | seed: int = 0, batch_size = None, acc_grad=1) -> None: |
| 251 | if num_replicas is None or rank is None or rank >= num_replicas or rank < 0: |
| 252 | raise ValueError( |
| 253 | f"Invalid num_replicas ({num_replicas}) or rank ({rank})") |
| 254 | assert batch_size is not None |
| 255 | self.batch_size = batch_size |
| 256 | |
| 257 | self.dataset = dataset |
| 258 | self.num_replicas = num_replicas |
| 259 | self.rank = rank |
| 260 | self.acc_grad = acc_grad |
| 261 | self.epoch = 0 |
| 262 | self.start_iter = 0 |
| 263 | |
| 264 | group_indices = dataset.groups() |
| 265 | global_bsz = batch_size * num_replicas * acc_grad |
| 266 | len_groups = [len(_) // global_bsz * global_bsz for _ in group_indices] |
| 267 | group_indices = [indices[:len_indices] for indices, len_indices in zip(group_indices, len_groups)] |
| 268 | group_n_batch = [len(_)//batch_size for _ in group_indices] |
| 269 | assert all([_%num_replicas==0 for _ in group_n_batch]) |
| 270 | n_total_batch = sum(group_n_batch) |
| 271 | |
| 272 | assert n_total_batch % self.num_replicas == 0 |
| 273 | |
| 274 | self.group_indices = group_indices |
| 275 | |
| 276 | self.total_size = n_total_batch * batch_size |
| 277 | self.num_samples = self.total_size // num_replicas |
| 278 | self.shuffle = shuffle |
| 279 | self.seed = seed |
| 280 | |
| 281 | def __iter__(self) -> Iterator: |
| 282 | global_batch_size = self.batch_size * self.num_replicas * self.acc_grad |
| 283 | if self.shuffle: |
| 284 | rng = np.random.default_rng(self.seed + self.epoch) |
| 285 | # self.group_indices should not be changed during shuffle. Only change copy. |
| 286 | group_indices_shuffle = copy.deepcopy(self.group_indices) |
| 287 | # for _ in group_indices_shuffle: |
| 288 | # rng.shuffle(_) |
| 289 | global_batched_indices = [ |
| 290 | indices_in_group[i:i+global_batch_size] |
| 291 | for indices_in_group in group_indices_shuffle |
| 292 | for i in range(0, len(indices_in_group), global_batch_size)] |
| 293 | rng.shuffle(global_batched_indices) |
| 294 | indices = [_ for batch_indices in global_batched_indices for _ in batch_indices] |
| 295 | else: |
| 296 | group_indices = copy.deepcopy(self.group_indices) |
| 297 | indices = [_ for indices_in_group in group_indices for _ in indices_in_group] |
| 298 | |
| 299 | assert len(indices) == self.total_size |
| 300 | |
| 301 | own_indices = [] |
| 302 | for start_pos in range(self.rank * self.batch_size, len(indices), self.num_replicas * self.batch_size): |
| 303 | own_indices += indices[start_pos: start_pos + self.batch_size] |