| import torch |
| import json |
|
|
| |
| def get_train_sampler(dataset, rank, world_size, global_batch_size, max_steps, |
| resume_step, seed): |
| sample_indices = torch.empty([max_steps * global_batch_size // world_size], |
| dtype=torch.long) |
| epoch_id, fill_ptr, offs = 0, 0, 0 |
| while fill_ptr < sample_indices.size(0): |
| g = torch.Generator() |
| g.manual_seed(seed + epoch_id) |
| epoch_sample_indices = torch.randperm(len(dataset), generator=g) |
| epoch_id += 1 |
| epoch_sample_indices = epoch_sample_indices[ |
| (rank + offs) % world_size::world_size |
| ] |
| offs = (offs + world_size - len(dataset) % world_size) % world_size |
| epoch_sample_indices = epoch_sample_indices[ |
| :sample_indices.size(0) - fill_ptr |
| ] |
| sample_indices[fill_ptr: fill_ptr + epoch_sample_indices.size(0)] = \ |
| epoch_sample_indices |
| fill_ptr += epoch_sample_indices.size(0) |
| return sample_indices[resume_step * global_batch_size // world_size:].tolist() |
|
|
|
|
|
|
|
|
| def get_packed_batch_sampler( |
| dataset, rank, world_size, max_steps, resume_step, seed |
| ): |
| sample_indices = [None for _ in range(max_steps)] |
| epoch_id, fill_ptr, offs = 0, 0, 0 |
| while fill_ptr < len(sample_indices): |
| g = torch.Generator() |
| g.manual_seed(seed + epoch_id) |
| epoch_sample_indices = torch.randperm(len(dataset), generator=g) |
| epoch_id += 1 |
| epoch_sample_indices = epoch_sample_indices[ |
| (rank + offs) % world_size::world_size |
| ] |
| offs = (offs + world_size - len(dataset) % world_size) % world_size |
| epoch_sample_indices = epoch_sample_indices[ |
| :len(sample_indices) - fill_ptr |
| ] |
| sample_indices[fill_ptr: fill_ptr + epoch_sample_indices.size(0)] = [ |
| dataset[i] for i in epoch_sample_indices |
| ] |
| fill_ptr += epoch_sample_indices.size(0) |
| return sample_indices[resume_step:] |
|
|
|
|
