| 51 | |
| 52 | |
| 53 | def get_train_sampler(dataset, rank, world_size, global_batch_size, max_steps, |
| 54 | resume_step, seed): |
| 55 | sample_indices = torch.empty([max_steps * global_batch_size // world_size], |
| 56 | dtype=torch.long) |
| 57 | epoch_id, fill_ptr, offs = 0, 0, 0 |
| 58 | while fill_ptr < sample_indices.size(0): |
| 59 | g = torch.Generator() |
| 60 | g.manual_seed(seed + epoch_id) |
| 61 | epoch_sample_indices = torch.randperm(len(dataset), generator=g) |
| 62 | epoch_id += 1 |
| 63 | epoch_sample_indices = epoch_sample_indices[ |
| 64 | (rank + offs) % world_size::world_size |
| 65 | ] |
| 66 | offs = (offs + world_size - len(dataset) % world_size) % world_size |
| 67 | epoch_sample_indices = epoch_sample_indices[ |
| 68 | :sample_indices.size(0) - fill_ptr |
| 69 | ] |
| 70 | sample_indices[fill_ptr: fill_ptr + epoch_sample_indices.size(0)] = \ |
| 71 | epoch_sample_indices |
| 72 | fill_ptr += epoch_sample_indices.size(0) |
| 73 | return sample_indices[resume_step * global_batch_size // world_size:].tolist() |
| 74 | |
| 75 | |
| 76 | @torch.no_grad() |