wehub-resource-sync
70 lines
2.6 KiB
Python
70 lines
2.6 KiB
Python
# coding: utf-8
|
|
|
|
import os
|
|
import random
|
|
import torch
|
|
from common.utils.logging import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
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):
|
|
# Process granularity: one shard per rank
|
|
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]
|
|
|
|
# ================== Add this log line ==================
|
|
if self.data_paths_per_rank and self.local_rank == 0: # Ensure the list is non-empty and log only on rank 0
|
|
logger.info(f"[Rank-Split-Check] Rank {self.local_rank} got {len(self.data_paths_per_rank)} files. "
|
|
f"First file: {os.path.basename(self.data_paths_per_rank[0])}")
|
|
# =======================================================
|
|
|
|
def get_data_paths_per_worker(self):
|
|
# Worker granularity: one shard per worker process
|
|
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]
|
|
|
|
# NOTE: The reverse order ::-1 is probably unnecessary
|
|
return data_paths_per_worker, worker_id
|
|
|
|
def __iter__(self):
|
|
raise NotImplementedError
|