wehub-resource-sync
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
from megatron.core import parallel_state
|
|
import torch
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
from general_util.logger import get_child_logger
|
|
|
|
logger = get_child_logger(__name__)
|
|
|
|
|
|
def get_model_parallel_group():
|
|
return parallel_state.get_tensor_model_parallel_group()
|
|
|
|
|
|
def get_model_parallel_rank():
|
|
return parallel_state.get_tensor_model_parallel_rank()
|
|
|
|
|
|
def get_model_parallel_world_size():
|
|
return parallel_state.get_tensor_model_parallel_world_size()
|
|
|
|
|
|
def get_data_parallel_group():
|
|
return parallel_state.get_data_parallel_group()
|
|
|
|
|
|
def get_data_parallel_rank():
|
|
return parallel_state.get_data_parallel_rank()
|
|
|
|
|
|
def get_data_parallel_world_size():
|
|
return parallel_state.get_data_parallel_world_size()
|
|
|
|
|
|
def prepare_distributed_sampler(dataset: torch.utils.data.Dataset, random_seed: int = 42, shuffle: bool = True):
|
|
if parallel_state.model_parallel_is_initialized():
|
|
sub_train_sampler = DistributedSampler(dataset,
|
|
shuffle=shuffle,
|
|
num_replicas=parallel_state.get_data_parallel_world_size(),
|
|
rank=parallel_state.get_data_parallel_rank(),
|
|
seed=random_seed)
|
|
else:
|
|
sub_train_sampler = DistributedSampler(dataset, shuffle=shuffle)
|
|
|
|
logger.info(f"Distributed Shuffling: {shuffle}")
|
|
return sub_train_sampler
|