wehub-resource-sync
124 lines
4.5 KiB
Python
124 lines
4.5 KiB
Python
import datetime
|
|
import os
|
|
import subprocess
|
|
|
|
import fairscale.nn.model_parallel.initialize as mpu
|
|
import torch
|
|
import torch.distributed as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
|
|
from omegaconf import DictConfig
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
from general_util.logger import get_child_logger
|
|
|
|
|
|
logger = get_child_logger(__name__)
|
|
|
|
|
|
def vanilla_torch_dist(cfg: DictConfig, backend="nccl"):
|
|
if "LOCAL_RANK" in os.environ and os.environ["LOCAL_RANK"] not in [-1, "-1"]:
|
|
cfg.local_rank = int(os.environ["LOCAL_RANK"])
|
|
|
|
if cfg.local_rank == -1 or cfg.no_cuda:
|
|
device = str(torch.device("cuda" if torch.cuda.is_available() and not cfg.no_cuda else "cpu"))
|
|
cfg.n_gpu = torch.cuda.device_count()
|
|
else: # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
|
|
torch.cuda.set_device(cfg.local_rank)
|
|
device = str(torch.device("cuda", cfg.local_rank))
|
|
dist.init_process_group(backend=backend, timeout=datetime.timedelta(seconds=7200))
|
|
cfg.n_gpu = 1
|
|
cfg.world_size = dist.get_world_size()
|
|
cfg.device = device
|
|
|
|
|
|
def setup_slurm_distributed(cfg: DictConfig, backend="nccl", port=None):
|
|
"""
|
|
Most code are copied from https://github.com/BIGBALLON/distribuuuu/blob/master/tutorial/mnmc_ddp_slurm.py.
|
|
"""
|
|
num_gpus = torch.cuda.device_count()
|
|
print(num_gpus)
|
|
if num_gpus <= 1 or cfg.no_cuda:
|
|
cfg.local_rank = -1
|
|
cfg.device = str(torch.device("cuda" if torch.cuda.is_available() and not cfg.no_cuda else "cpu"))
|
|
cfg.n_gpu = min(num_gpus, 1)
|
|
cfg.ddp_eval = False
|
|
return
|
|
|
|
# Data Parallel or Model Parallel on multiple GPUs with single task.
|
|
if int(os.environ["SLURM_NTASKS"]) == 1:
|
|
cfg.n_gpu = num_gpus
|
|
cfg.ddp_eval = False
|
|
cfg.device = str(torch.device("cuda"))
|
|
cfg.local_rank = -1
|
|
return
|
|
|
|
proc_id = int(os.environ["SLURM_PROCID"])
|
|
n_tasks = int(os.environ["SLURM_NTASKS"])
|
|
node_list = os.environ["SLURM_NODELIST"]
|
|
|
|
torch.cuda.set_device(proc_id % num_gpus)
|
|
|
|
addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1")
|
|
# specify master port
|
|
if port is not None:
|
|
os.environ["MASTER_PORT"] = str(port)
|
|
elif "MASTER_PORT" not in os.environ:
|
|
os.environ["MASTER_PORT"] = "29500"
|
|
if "MASTER_ADDR" not in os.environ:
|
|
os.environ["MASTER_ADDR"] = addr
|
|
|
|
os.environ["WORLD_SIZE"] = str(n_tasks)
|
|
os.environ["LOCAL_RANK"] = str(proc_id % num_gpus)
|
|
os.environ["RANK"] = str(proc_id)
|
|
|
|
cfg.n_gpu = 1
|
|
cfg.local_rank = int(os.environ["LOCAL_RANK"])
|
|
# cfg.local_rank = int(os.environ["RANK"])
|
|
cfg.world_size = int(os.environ["WORLD_SIZE"])
|
|
cfg.device = str(torch.device("cuda", cfg.local_rank))
|
|
|
|
dist.init_process_group(backend=backend, world_size=int(os.environ["WORLD_SIZE"]), rank=int(os.environ["RANK"]))
|
|
|
|
# print(cfg.n_gpu, cfg.local_rank, cfg.world_size, cfg.device)
|
|
# print(cfg.local_rank)
|
|
cfg.local_rank = dist.get_rank()
|
|
# print(cfg.local_rank)
|
|
|
|
|
|
def print_rank_0(msg, rank=0):
|
|
if rank <= 0:
|
|
print(msg)
|
|
|
|
|
|
def print_all_ranks(tag, value, rank):
|
|
world_size = dist.get_world_size()
|
|
all_tensor = torch.zeros(world_size, dtype=torch.float32).to(get_accelerator().current_device_name())
|
|
all_tensor[rank] = value
|
|
dist.all_reduce(all_tensor, op=dist.ReduceOp.SUM)
|
|
print_rank_0(f'{tag} {all_tensor}', rank)
|
|
|
|
|
|
def get_pipeline_parallel_world_size() -> int:
|
|
"""Return world size for the model parallel group."""
|
|
return torch.distributed.get_world_size(group=get_pipeline_parallel_group())
|
|
|
|
|
|
def get_pipeline_parallel_rank() -> int:
|
|
"""Return my rank for the model parallel group."""
|
|
return torch.distributed.get_rank(group=get_pipeline_parallel_group())
|
|
|
|
|
|
def prepare_distributed_sampler(dataset: torch.utils.data.Dataset, random_seed: int = 42, shuffle: bool = True):
|
|
if mpu.model_parallel_is_initialized():
|
|
sub_train_sampler = DistributedSampler(dataset,
|
|
shuffle=shuffle,
|
|
num_replicas=mpu.get_data_parallel_world_size(),
|
|
rank=mpu.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
|