|
|
|
"""Contains utility functions used for distribution.""" |
|
|
|
import contextlib |
|
import os |
|
import subprocess |
|
|
|
import torch |
|
import torch.distributed as dist |
|
import torch.multiprocessing as mp |
|
|
|
__all__ = ['init_dist', 'exit_dist', 'ddp_sync', 'get_ddp_module'] |
|
|
|
|
|
def init_dist(launcher, backend='nccl', **kwargs): |
|
"""Initializes distributed environment.""" |
|
if mp.get_start_method(allow_none=True) is None: |
|
mp.set_start_method('spawn') |
|
if launcher == 'pytorch': |
|
rank = int(os.environ['RANK']) |
|
num_gpus = torch.cuda.device_count() |
|
torch.cuda.set_device(rank % num_gpus) |
|
dist.init_process_group(backend=backend, **kwargs) |
|
elif launcher == 'slurm': |
|
proc_id = int(os.environ['SLURM_PROCID']) |
|
ntasks = int(os.environ['SLURM_NTASKS']) |
|
node_list = os.environ['SLURM_NODELIST'] |
|
num_gpus = torch.cuda.device_count() |
|
torch.cuda.set_device(proc_id % num_gpus) |
|
addr = subprocess.getoutput( |
|
f'scontrol show hostname {node_list} | head -n1') |
|
port = os.environ.get('PORT', 29500) |
|
os.environ['MASTER_PORT'] = str(port) |
|
os.environ['MASTER_ADDR'] = addr |
|
os.environ['WORLD_SIZE'] = str(ntasks) |
|
os.environ['RANK'] = str(proc_id) |
|
dist.init_process_group(backend=backend) |
|
else: |
|
raise NotImplementedError(f'Not implemented launcher type: ' |
|
f'`{launcher}`!') |
|
|
|
|
|
def exit_dist(): |
|
"""Exits the distributed environment.""" |
|
if dist.is_initialized(): |
|
dist.destroy_process_group() |
|
|
|
|
|
@contextlib.contextmanager |
|
def ddp_sync(model, sync): |
|
"""Controls whether the `DistributedDataParallel` model should be synced.""" |
|
assert isinstance(model, torch.nn.Module) |
|
is_ddp = isinstance(model, torch.nn.parallel.DistributedDataParallel) |
|
if sync or not is_ddp: |
|
yield |
|
else: |
|
with model.no_sync(): |
|
yield |
|
|
|
|
|
def get_ddp_module(model): |
|
"""Gets the module from `DistributedDataParallel`.""" |
|
assert isinstance(model, torch.nn.Module) |
|
is_ddp = isinstance(model, torch.nn.parallel.DistributedDataParallel) |
|
if is_ddp: |
|
return model.module |
|
return model |
|
|
