Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright 2025 Bytedance Ltd. and/or its affiliates. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import random | |
| import torch | |
| class DistributedIterableDataset(torch.utils.data.IterableDataset): | |
| def __init__(self, dataset_name, local_rank=0, world_size=1, num_workers=8): | |
| self.dataset_name = dataset_name | |
| self.local_rank = local_rank | |
| self.world_size = world_size | |
| self.num_workers = num_workers | |
| self.rng = random.Random() | |
| self.data_paths = None | |
| def get_data_paths(self, *args, **kwargs): | |
| raise NotImplementedError | |
| def set_epoch(self, seed=42): | |
| if self.data_paths is None: | |
| return | |
| if isinstance(self.data_paths[0], tuple): | |
| data_paths = sorted(self.data_paths, key=lambda x: (x[0], x[1])) | |
| elif isinstance(self.data_paths[0], str): | |
| data_paths = sorted(self.data_paths) | |
| else: | |
| raise ValueError(f"Unknown data_paths type: {type(self.data_paths[0])}") | |
| self.rng.seed(seed) | |
| self.rng.shuffle(data_paths) | |
| num_files_per_rank = len(data_paths) // self.world_size | |
| local_start = self.local_rank * num_files_per_rank | |
| local_end = (self.local_rank + 1) * num_files_per_rank | |
| self.num_files_per_rank = num_files_per_rank | |
| self.data_paths_per_rank = data_paths[local_start:local_end] | |
| def get_data_paths_per_worker(self): | |
| if self.data_paths is None: | |
| return None | |
| info = torch.utils.data.get_worker_info() | |
| if info is None: | |
| # Single worker: Use all files assigned to the rank | |
| return self.data_paths_per_rank, 0 | |
| worker_id = info.id | |
| num_files_per_worker = self.num_files_per_rank // info.num_workers | |
| start = num_files_per_worker * worker_id | |
| end = num_files_per_worker * (worker_id + 1) | |
| data_paths_per_worker = self.data_paths_per_rank[start:end] | |
| return data_paths_per_worker[::-1], worker_id | |
| def __iter__(self): | |
| raise NotImplementedError | |