812 lines
28 KiB
Python
812 lines
28 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import functools
|
|
import os
|
|
import pickle
|
|
import random
|
|
import socket
|
|
import warnings
|
|
from collections import OrderedDict
|
|
from datetime import timedelta
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.autograd import Function
|
|
|
|
from scepter.modules.utils.model import StdMsg
|
|
|
|
__all__ = [
|
|
'gather_data', 'we', 'broadcast', 'barrier', 'reduce_scatter', 'reduce',
|
|
'all_reduce', 'send', 'recv', 'isend', 'irecv', 'scatter',
|
|
'shared_random_seed'
|
|
]
|
|
|
|
try:
|
|
from onnxruntime.transformers.benchmark_helper import set_random_seed
|
|
except Exception:
|
|
|
|
def set_random_seed(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
|
|
def get_dist_info():
|
|
if dist.is_available() and dist.is_initialized():
|
|
return dist.get_rank(), dist.get_world_size()
|
|
else:
|
|
return 0, 1
|
|
|
|
|
|
def gather_data(data):
|
|
""" Gather tensors and other picklable objects to rank 0.
|
|
Will recursively walk through inner list and dict values.
|
|
|
|
Args:
|
|
data (any): Anything.
|
|
|
|
Returns:
|
|
A object has same structure with input `data`.
|
|
"""
|
|
if not we.is_distributed:
|
|
return data
|
|
if isinstance(data, torch.Tensor):
|
|
return gather_gpu_tensors(data)
|
|
elif isinstance(data, dict):
|
|
# Keep in order, dict type DO NOT guarantee a fixed key order
|
|
keys = sorted(list(data.keys()))
|
|
ret = OrderedDict()
|
|
for key in keys:
|
|
ret[key] = gather_data(data[key])
|
|
return ret
|
|
elif isinstance(data, list):
|
|
return gather_list(data)
|
|
else:
|
|
return gather_picklable(data)
|
|
|
|
|
|
def gather_list(data):
|
|
""" Gather list of picklable objects to a new list on rank 0.
|
|
Will NOT recursively walk through.
|
|
|
|
Args:
|
|
data (list): List of picklable things.
|
|
|
|
Returns:
|
|
A new flat list.
|
|
"""
|
|
rank, _ = get_dist_info()
|
|
list_of_list = gather_picklable(data)
|
|
if rank == 0:
|
|
return sum(list_of_list, [])
|
|
|
|
|
|
def gather_picklable(data):
|
|
""" Gather picklable object to a list on rank 0.
|
|
Will NOT recursively walk through.
|
|
|
|
Args:
|
|
data (picklable): Picklable data.
|
|
|
|
Returns:
|
|
A list contains data collected.
|
|
"""
|
|
from packaging import version
|
|
from torch.version import __version__
|
|
if version.parse(__version__) < version.parse('1.8.0'):
|
|
return _gather_picklable_custom(data)
|
|
else:
|
|
rank, world_size = we.rank, we.world_size
|
|
obj_list = [None for _ in range(world_size)]
|
|
dist.all_gather_object(obj_list, data)
|
|
if rank == 0:
|
|
return obj_list
|
|
|
|
|
|
def _gather_picklable_custom(data):
|
|
""" Custom implementation function to gather picklable object to a list on rank 0.
|
|
If torch version is lower than 1.8.0, use this.
|
|
|
|
Args:
|
|
data (picklable): Picklable data.
|
|
|
|
Returns:
|
|
A list contains data collected.
|
|
"""
|
|
import pickle
|
|
byte_tensor = torch.tensor(bytearray(pickle.dumps(data)),
|
|
dtype=torch.uint8,
|
|
device='cuda')
|
|
rank, world_size = we.rank, we.world_size
|
|
shape_tensor = torch.tensor(byte_tensor.shape, device='cuda')
|
|
shape_list = [shape_tensor.clone() for _ in range(world_size)]
|
|
dist.all_gather(shape_list, shape_tensor)
|
|
shape_max = torch.tensor(shape_list).max()
|
|
|
|
tensor_send = torch.zeros(shape_max,
|
|
dtype=byte_tensor.dtype,
|
|
device='cuda')
|
|
tensor_send[0:shape_tensor[0]] = byte_tensor
|
|
tensor_list = [torch.zeros_like(tensor_send) for _ in range(world_size)]
|
|
dist.all_gather(tensor_list, tensor_send)
|
|
|
|
if rank == 0:
|
|
data_out = []
|
|
for tensor_recv, shape_recv in zip(tensor_list, shape_list):
|
|
new_data = pickle.loads(
|
|
tensor_recv[:shape_recv[0]].cpu().numpy().tobytes())
|
|
data_out.append(new_data)
|
|
return data_out
|
|
|
|
|
|
def gather_gpu_tensors(tensor, all_recv=False, is_cat=True):
|
|
"""
|
|
Args:
|
|
tensor (torch.Tensor):
|
|
all_recv: Gather tensor to rank 0 and concat it.
|
|
|
|
Returns:
|
|
A new tensor.
|
|
"""
|
|
assert dist.get_backend() == 'nccl'
|
|
|
|
device = tensor.device
|
|
if device.type == 'cpu':
|
|
tensor = tensor.to(we.device_id)
|
|
|
|
rank, world_size = we.rank, we.world_size
|
|
|
|
shape_tensor = torch.tensor(tensor.shape[0], device='cuda')
|
|
shape_list = [shape_tensor.clone() for _ in range(world_size)]
|
|
dist.all_gather(shape_list, shape_tensor)
|
|
shape_max = torch.tensor(shape_list).max()
|
|
|
|
tensor_send = torch.zeros((shape_max, *tensor.shape[1:]),
|
|
dtype=tensor.dtype,
|
|
device='cuda')
|
|
tensor_send[0:tensor.shape[0]] = tensor
|
|
tensor_list = [torch.zeros_like(tensor_send) for _ in range(world_size)]
|
|
dist.all_gather(tensor_list, tensor_send)
|
|
if not all_recv:
|
|
if rank == 0:
|
|
if not is_cat:
|
|
return tensor_list, shape_list
|
|
tensors_out = []
|
|
for tensor_recv, shape_recv in zip(tensor_list, shape_list):
|
|
tensors_out.append(tensor_recv[0:shape_recv])
|
|
tensor_out = torch.cat(tensors_out).contiguous()
|
|
if device.type == 'cpu':
|
|
tensor_out = tensor_out.cpu()
|
|
del tensor_list, shape_list
|
|
return tensor_out
|
|
else:
|
|
del tensor_list, shape_list
|
|
else:
|
|
if not is_cat:
|
|
return tensor_list, shape_list
|
|
tensors_out = []
|
|
for tensor_recv, shape_recv in zip(tensor_list, shape_list):
|
|
tensors_out.append(tensor_recv[0:shape_recv])
|
|
tensor_out = torch.cat(tensors_out).contiguous()
|
|
if device.type == 'cpu':
|
|
tensor_out = tensor_out.cpu()
|
|
del tensor_list, shape_list
|
|
return tensor_out
|
|
|
|
|
|
def broadcast(tensor, src, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
return dist.broadcast(tensor, src, group, **kwargs)
|
|
|
|
|
|
def broadcast_object_list(object_list, src, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
return dist.broadcast_object_list(object_list, src, group, **kwargs)
|
|
|
|
|
|
def barrier():
|
|
if we.is_distributed:
|
|
dist.barrier()
|
|
|
|
|
|
def all_gather(tensor, uniform_size=True, group=None, **kwargs):
|
|
world_size = dist.get_world_size(group)
|
|
if world_size == 1:
|
|
return [tensor]
|
|
assert tensor.is_contiguous(), \
|
|
'ops.all_gather requires the tensor to be contiguous()'
|
|
|
|
if uniform_size:
|
|
tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
|
|
dist.all_gather(tensor_list, tensor, group, **kwargs)
|
|
return tensor_list
|
|
else:
|
|
# collect tensor shapes across GPUs
|
|
shape = tuple(tensor.shape)
|
|
shape_list = generalized_all_gather(shape, group)
|
|
|
|
# flatten the tensor
|
|
tensor = tensor.reshape(-1)
|
|
size = int(np.prod(shape))
|
|
size_list = [int(np.prod(u)) for u in shape_list]
|
|
max_size = max(size_list)
|
|
|
|
# pad to maximum size
|
|
if size != max_size:
|
|
padding = tensor.new_zeros(max_size - size)
|
|
tensor = torch.cat([tensor, padding], dim=0)
|
|
|
|
# all_gather
|
|
tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
|
|
dist.all_gather(tensor_list, tensor, group, **kwargs)
|
|
|
|
# reshape tensors
|
|
tensor_list = [
|
|
t[:n].view(s)
|
|
for t, n, s in zip(tensor_list, size_list, shape_list)
|
|
]
|
|
return tensor_list
|
|
|
|
|
|
def _pad_to_largest_tensor(tensor, group):
|
|
world_size = dist.get_world_size(group=group)
|
|
assert world_size >= 1, \
|
|
'gather/all_gather must be called from ranks within' \
|
|
'the give group!'
|
|
local_size = torch.tensor([tensor.numel()],
|
|
dtype=torch.int64,
|
|
device=tensor.device)
|
|
size_list = [
|
|
torch.zeros([1], dtype=torch.int64, device=tensor.device)
|
|
for _ in range(world_size)
|
|
]
|
|
|
|
# gather tensors and compute the maximum size
|
|
dist.all_gather(size_list, local_size, group=group)
|
|
size_list = [int(size.item()) for size in size_list]
|
|
max_size = max(size_list)
|
|
|
|
# pad tensors to the same size
|
|
if local_size != max_size:
|
|
padding = torch.zeros((max_size - local_size, ),
|
|
dtype=torch.uint8,
|
|
device=tensor.device)
|
|
tensor = torch.cat((tensor, padding), dim=0)
|
|
return size_list, tensor
|
|
|
|
|
|
def generalized_all_gather(data, group=None):
|
|
if dist.get_world_size(group) == 1:
|
|
return [data]
|
|
if group is None:
|
|
group = get_global_gloo_group()
|
|
|
|
tensor = _serialize_to_tensor(data, group)
|
|
size_list, tensor = _pad_to_largest_tensor(tensor, group)
|
|
max_size = max(size_list)
|
|
|
|
# receiving tensors from all ranks
|
|
tensor_list = [
|
|
torch.empty((max_size, ), dtype=torch.uint8, device=tensor.device)
|
|
for _ in size_list
|
|
]
|
|
dist.all_gather(tensor_list, tensor, group=group)
|
|
|
|
data_list = []
|
|
for size, tensor in zip(size_list, tensor_list):
|
|
buffer = tensor.cpu().numpy().tobytes()[:size]
|
|
data_list.append(pickle.loads(buffer))
|
|
return data_list
|
|
|
|
|
|
@functools.lru_cache()
|
|
def get_global_gloo_group():
|
|
backend = dist.get_backend()
|
|
assert backend in ['gloo', 'nccl']
|
|
if backend == 'nccl':
|
|
return dist.new_group(backend='gloo')
|
|
else:
|
|
return dist.group.WORLD
|
|
|
|
|
|
def reduce_scatter(output,
|
|
input_list,
|
|
op=dist.ReduceOp.SUM,
|
|
group=None,
|
|
**kwargs):
|
|
if we.is_distributed:
|
|
return dist.reduce_scatter(output, input_list, op, group, **kwargs)
|
|
|
|
|
|
def all_reduce(tensor, op=dist.ReduceOp.SUM, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
dist.all_reduce(tensor, op, group, **kwargs)
|
|
return tensor
|
|
|
|
|
|
def reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
dist.reduce(tensor, dst, op, group, **kwargs)
|
|
return tensor
|
|
|
|
|
|
def _serialize_to_tensor(data):
|
|
buffer = pickle.dumps(data)
|
|
storage = torch.ByteStorage.from_buffer(buffer)
|
|
tensor = torch.ByteTensor(storage)
|
|
return tensor
|
|
|
|
|
|
def _unserialize_from_tensor(recv_data):
|
|
buffer = recv_data.cpu().numpy().tobytes()
|
|
return pickle.loads(buffer)
|
|
|
|
|
|
def find_free_port():
|
|
# Copied from https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py # noqa: E501
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
# Binding to port 0 will cause the OS to find an available port for us
|
|
sock.bind(('', 0))
|
|
port = sock.getsockname()[1]
|
|
sock.close()
|
|
# NOTE: there is still a chance the port could be taken by other processes.
|
|
return port
|
|
|
|
|
|
def is_free_port(port):
|
|
ips = socket.gethostbyname_ex(socket.gethostname())[-1]
|
|
ips.append('localhost')
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
return all(s.connect_ex((ip, port)) != 0 for ip in ips)
|
|
|
|
|
|
def send(tensor, dst, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
assert tensor.is_contiguous(
|
|
), 'ops.send requires the tensor to be contiguous()'
|
|
return dist.send(tensor, dst, group, **kwargs)
|
|
|
|
|
|
def recv(tensor, src=None, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
assert tensor.is_contiguous(
|
|
), 'ops.recv requires the tensor to be contiguous()'
|
|
return dist.recv(tensor, src, group, **kwargs)
|
|
|
|
|
|
def isend(tensor, dst, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
assert tensor.is_contiguous(
|
|
), 'ops.isend requires the tensor to be contiguous()'
|
|
return dist.isend(tensor, dst, group, **kwargs)
|
|
|
|
|
|
def irecv(tensor, src=None, group=None, **kwargs):
|
|
if we.is_distributed:
|
|
assert tensor.is_contiguous(
|
|
), 'ops.irecv requires the tensor to be contiguous()'
|
|
return dist.irecv(tensor, src, group, **kwargs)
|
|
|
|
|
|
def scatter(data, scatter_list=None, src=0, group=None, **kwargs):
|
|
r"""NOTE: only supports CPU tensor communication.
|
|
"""
|
|
world_size = we.world_size
|
|
if world_size == 1:
|
|
data.copy_(scatter_list[0])
|
|
if group is None:
|
|
group = get_global_gloo_group()
|
|
return dist.scatter(data, scatter_list, src, group, **kwargs)
|
|
|
|
|
|
def shared_random_seed():
|
|
seed = np.random.randint(2**31)
|
|
all_seeds, _ = gather_gpu_tensors(seed, all_recv=True, is_cat=False)
|
|
return all_seeds[0]
|
|
|
|
|
|
def all_to_all(x, scatter_dim, gather_dim, group=None, **kwargs):
|
|
"""
|
|
`scatter` along one dimension and `gather` along another.
|
|
"""
|
|
world_size = dist.get_world_size(group) if we.is_distributed else 1
|
|
if world_size > 1:
|
|
inputs = [u.contiguous() for u in x.chunk(world_size, dim=scatter_dim)]
|
|
outputs = [torch.empty_like(u) for u in inputs]
|
|
dist.all_to_all(outputs, inputs, group=group, **kwargs)
|
|
x = torch.cat(outputs, dim=gather_dim).contiguous()
|
|
return x
|
|
|
|
|
|
def _split(input, dim, group):
|
|
# skip if world_size == 1
|
|
rank = dist.get_rank(group=group)
|
|
world_size = dist.get_world_size(group=group)
|
|
if world_size == 1:
|
|
return input
|
|
|
|
# split sequence
|
|
assert input.size(dim) % world_size == 0
|
|
return input.chunk(world_size, dim=dim)[rank].contiguous()
|
|
|
|
|
|
def _gather(input, dim, group):
|
|
# skip if world_size == 1
|
|
world_size = dist.get_world_size(group=group)
|
|
if world_size == 1:
|
|
return input
|
|
|
|
# gather sequence
|
|
output = all_gather(input, uniform_size=True, group=group)
|
|
return torch.cat(output, dim=dim).contiguous()
|
|
|
|
|
|
class AllToAll(Function):
|
|
@staticmethod
|
|
def forward(ctx, input, scatter_dim, gather_dim, group):
|
|
ctx.scatter_dim = scatter_dim
|
|
ctx.gather_dim = gather_dim
|
|
ctx.group = group
|
|
return all_to_all(input, scatter_dim, gather_dim, group)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
return (all_to_all(grad_output, ctx.gather_dim, ctx.scatter_dim,
|
|
ctx.group), None, None, None)
|
|
|
|
|
|
class GradScaler(Function):
|
|
@staticmethod
|
|
def forward(ctx, input, scale):
|
|
ctx.scale = scale
|
|
return input
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
if ctx.scale != 1:
|
|
grad_output = grad_output * ctx.scale
|
|
return grad_output, None
|
|
|
|
|
|
class AllGather(Function):
|
|
@staticmethod
|
|
def forward(ctx, input, dim, group=None):
|
|
ctx.dim = dim
|
|
ctx.group = group
|
|
output = all_gather(input, uniform_size=True, group=group)
|
|
return torch.cat(output, dim=dim).contiguous()
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
rank = dist.get_rank(group=ctx.group)
|
|
world_size = dist.get_world_size(group=ctx.group)
|
|
return grad_output.chunk(world_size,
|
|
dim=ctx.dim)[rank].contiguous(), None, None
|
|
|
|
|
|
def diff_all_to_all(input, scatter_dim, gather_dim, group=None):
|
|
return AllToAll.apply(input, scatter_dim, gather_dim, group)
|
|
|
|
|
|
def diff_scatter_sequence(input, dim, group=None):
|
|
rank = dist.get_rank(group)
|
|
world_size = dist.get_world_size(group)
|
|
output = input.chunk(world_size, dim=dim)[rank].contiguous()
|
|
return GradScaler.apply(output, 1. / world_size)
|
|
|
|
|
|
def diff_gather_sequence(input, dim, group=None):
|
|
world_size = dist.get_world_size(group)
|
|
output = AllGather.apply(input, dim, group)
|
|
return GradScaler.apply(output, world_size)
|
|
|
|
|
|
class SplitForwardGatherBackward(Function):
|
|
@staticmethod
|
|
def forward(ctx, input, dim, group=None, grad_scale=None):
|
|
ctx.dim = dim
|
|
ctx.group = group
|
|
ctx.grad_scale = grad_scale
|
|
return _split(input, dim, group)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
if ctx.grad_scale == 'up':
|
|
grad_output = grad_output * dist.get_world_size(group=ctx.group)
|
|
elif ctx.grad_scale == 'down':
|
|
grad_output = grad_output / dist.get_world_size(group=ctx.group)
|
|
return _gather(grad_output, ctx.dim, ctx.group), None, None, None
|
|
|
|
|
|
class GatherForwardSplitBackward(Function):
|
|
@staticmethod
|
|
def forward(ctx, input, dim, group=None, grad_scale=None):
|
|
ctx.dim = dim
|
|
ctx.group = group
|
|
ctx.grad_scale = grad_scale
|
|
return _gather(input, dim, group)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
if ctx.grad_scale == 'up':
|
|
grad_output = grad_output * dist.get_world_size(group=ctx.group)
|
|
elif ctx.grad_scale == 'down':
|
|
grad_output = grad_output / dist.get_world_size(group=ctx.group)
|
|
return _split(grad_output, ctx.dim, ctx.group), None, None, None
|
|
|
|
|
|
def split_forward_gather_backward(input, dim, group=None, grad_scale=None):
|
|
return SplitForwardGatherBackward.apply(input, dim, group, grad_scale)
|
|
|
|
|
|
def gather_forward_split_backward(input, dim, group=None, grad_scale=None):
|
|
return GatherForwardSplitBackward.apply(input, dim, group, grad_scale)
|
|
|
|
|
|
global we
|
|
|
|
|
|
def mp_worker(gpu, ngpus_per_node, cfg, fn, pmi_rank, world_size, work_env):
|
|
rank = pmi_rank * ngpus_per_node + gpu
|
|
work_env.device_id = gpu % ngpus_per_node
|
|
work_env.rank = rank
|
|
dist.init_process_group(backend='nccl',
|
|
world_size=world_size,
|
|
rank=rank,
|
|
timeout=timedelta(seconds=18000))
|
|
torch.backends.cudnn.deterministic = cfg.ENV.get('CUDNN_DETERMINISTIC',
|
|
True)
|
|
torch.backends.cudnn.benchmark = cfg.ENV.get('CUDNN_BENCHMARK', False)
|
|
torch.cuda.set_device(work_env.device_id)
|
|
if work_env.logger is not None:
|
|
work_env.logger.info(
|
|
f'Now running in the distributed environment with world size {work_env.world_size}!'
|
|
)
|
|
work_env.logger.info(f'PMI rank {pmi_rank}!')
|
|
work_env.logger.info(f'Nums of gpu {ngpus_per_node}!')
|
|
work_env.logger.info(
|
|
f'Current rank {work_env.rank} current devices num {ngpus_per_node} \n'
|
|
f'current machine rank {pmi_rank} and all world size {world_size}')
|
|
# model parallel
|
|
tensor_parallel_size = cfg.ENV.get('TENSOR_PARALLEL_SIZE', 1)
|
|
pipeline_parallel_size = cfg.ENV.get('PIPELINE_PARALLEL_SIZE', 1)
|
|
|
|
env_dict = work_env.get_env()
|
|
|
|
if tensor_parallel_size * pipeline_parallel_size > 1:
|
|
'''
|
|
'''
|
|
assert world_size % tensor_parallel_size == 0
|
|
assert world_size % (tensor_parallel_size *
|
|
pipeline_parallel_size) == 0
|
|
data_parallel_size = world_size // (tensor_parallel_size *
|
|
pipeline_parallel_size)
|
|
mesh = torch.arange(world_size).view(data_parallel_size,
|
|
pipeline_parallel_size,
|
|
tensor_parallel_size)
|
|
index = torch.where(mesh == rank)
|
|
assert all(u.numel() == 1 for u in index)
|
|
index = [u.item() for u in index]
|
|
for j in range(pipeline_parallel_size):
|
|
for k in range(tensor_parallel_size):
|
|
group = dist.new_group(mesh[:, j, k].tolist())
|
|
if j == index[1] and k == index[2]:
|
|
env_dict['data_parallel_group'] = group
|
|
for i in range(data_parallel_size):
|
|
for j in range(pipeline_parallel_size):
|
|
group = dist.new_group(mesh[i, j, :].tolist())
|
|
if i == index[0] and j == index[1]:
|
|
env_dict['tensor_parallel_group'] = group
|
|
for i in range(data_parallel_size):
|
|
for k in range(tensor_parallel_size):
|
|
ranks = mesh[i, :, k].tolist()
|
|
group = dist.new_group(ranks)
|
|
if i == index[0] and k == index[2]:
|
|
env_dict['pipeline_parallel_group'] = group
|
|
env_dict['pipeline_parallel_ranks'] = ranks
|
|
we.set_env(env_dict)
|
|
work_env.logger.info(str(we))
|
|
fn(cfg)
|
|
torch.cuda.synchronize()
|
|
barrier()
|
|
|
|
|
|
class Workenv(object):
|
|
def __init__(self):
|
|
self.initialized = False
|
|
self.is_distributed = False
|
|
self.sync_bn = False
|
|
self.rank = 0
|
|
self.world_size = 1
|
|
self.device_id = 0
|
|
self.backend = ''
|
|
self.device_count = 1
|
|
self.seed = 2023
|
|
self.debug = False
|
|
self.use_pl = False
|
|
self.launcher = 'spawn' if torch.cuda.device_count() > 1 else None
|
|
self.data_online = False
|
|
self.share_storage = False
|
|
|
|
self.data_parallel_group = None
|
|
self.tensor_parallel_group = None
|
|
self.pipleline_parallel_group = None
|
|
self.pipeline_parallel_ranks = None
|
|
|
|
def init_env(self, config, fn, logger=None):
|
|
# if use pytorch_lightning: then direct use pytorch_lightning.
|
|
config.ENV = config.get('ENV', {})
|
|
self.seed = config.ENV.get('SEED', 2023)
|
|
self.debug = os.environ.get('ES_DEBUG', None) == 'true'
|
|
|
|
if logger is None:
|
|
self.logger = StdMsg(name='env')
|
|
else:
|
|
self.logger = logger
|
|
|
|
self.sys_envs = config.ENV.get('SYS_ENVS', None)
|
|
if self.sys_envs:
|
|
for k, v in self.sys_envs.items():
|
|
os.environ[k] = v
|
|
self.logger.info(f'Set env variable {k}={v}')
|
|
set_random_seed(self.seed)
|
|
self.logger.info(f'And running with seed {self.seed}!')
|
|
|
|
if config.ENV.get('USE_PL', False):
|
|
self.use_pl = config.ENV.USE_PL
|
|
fn(config)
|
|
return
|
|
if hasattr(config, 'args') and hasattr(config.args, 'launcher'):
|
|
self.launcher = config.args.launcher
|
|
|
|
self.data_online = os.environ.get('DATA_ONLINE', None) == 'true'
|
|
self.share_storage = os.environ.get('SHARE_STORAGE', None) == 'true'
|
|
|
|
if not torch.cuda.is_available():
|
|
self.device_id = 'cpu'
|
|
fn(config)
|
|
return
|
|
|
|
if (os.environ.get('WORLD_SIZE') is None or int(os.environ.get('WORLD_SIZE')) == 1) \
|
|
and torch.cuda.device_count() == 1 and not self.launcher == 'dist':
|
|
self.device_id = 0
|
|
fn(config)
|
|
return
|
|
|
|
if self.launcher == 'torchrun':
|
|
try:
|
|
torch.multiprocessing.set_start_method('spawn')
|
|
except Exception as e:
|
|
warnings.warn(f'{e}')
|
|
# checking train mode is distributed or not
|
|
if not os.environ.get('WORLD_SIZE') is None:
|
|
if self.logger is not None:
|
|
self.logger.info(
|
|
f"Now running in the distributed environment with {os.environ.get('WORLD_SIZE')}!"
|
|
)
|
|
self.is_distributed = True
|
|
if not self.initialized:
|
|
if self.is_distributed:
|
|
self.backend = config.ENV.get('BACKEND', 'nccl')
|
|
self.sync_bn = config.ENV.get('SYNC_BN', False)
|
|
dist.init_process_group(backend=self.backend,
|
|
timeout=timedelta(seconds=18000))
|
|
# dist.barrier()
|
|
self.initialized = True
|
|
if dist.is_initialized():
|
|
self.rank, self.world_size = dist.get_rank(
|
|
), dist.get_world_size()
|
|
if self.logger is not None:
|
|
self.logger.info(f'And running in rank {self.rank}!')
|
|
self.logger.info(
|
|
f"And cuda visible devices {os.environ.get('CUDA_VISIBLE_DEVICES')}!"
|
|
)
|
|
else:
|
|
self.rank, self.world_size = 0, 1
|
|
local_devices = os.environ.get(
|
|
'LOCAL_WORLD_SIZE') or torch.cuda.device_count()
|
|
local_devices = int(local_devices)
|
|
self.device_count = local_devices
|
|
self.device_id = self.rank % local_devices
|
|
self.logger.info(f"We's attributes: \n"
|
|
f' launcher {self.launcher} \n'
|
|
f' rank {self.rank} \n'
|
|
f' world size {self.world_size} \n'
|
|
f' device_id {self.device_id}')
|
|
torch.cuda.set_device(self.device_id)
|
|
torch.backends.cudnn.deterministic = config.ENV.get(
|
|
'CUDNN_DETERMINISTIC', True)
|
|
torch.backends.cudnn.benchmark = config.ENV.get(
|
|
'CUDNN_BENCHMARK', False)
|
|
fn(config)
|
|
return
|
|
else:
|
|
import torch.multiprocessing as mp
|
|
if 'MASTER_ADDR' not in os.environ:
|
|
os.environ['MASTER_ADDR'] = 'localhost'
|
|
if 'MASTER_PORT' not in os.environ:
|
|
os.environ['MASTER_PORT'] = '14567'
|
|
pmi_rank = int(os.environ.get('RANK', 0))
|
|
pmi_world_size = int(os.environ.get('WORLD_SIZE', 1))
|
|
ngpus_per_node = os.environ.get(
|
|
'LOCAL_WORLD_SIZE') or torch.cuda.device_count()
|
|
ngpus_per_node = int(ngpus_per_node)
|
|
self.device_count = ngpus_per_node
|
|
world_size = ngpus_per_node * pmi_world_size
|
|
self.world_size = world_size
|
|
if self.world_size >= 1:
|
|
self.is_distributed = True
|
|
self.initialized = True
|
|
if self.is_distributed:
|
|
self.backend = config.ENV.get('BACKEND', 'nccl')
|
|
self.sync_bn = config.ENV.get('SYNC_BN', False)
|
|
spawn_join = config.ENV.get('SPAWN_JOIN', True)
|
|
context = mp.spawn(mp_worker,
|
|
nprocs=ngpus_per_node,
|
|
join=spawn_join,
|
|
args=(ngpus_per_node, config, fn, pmi_rank, world_size,
|
|
self))
|
|
return context
|
|
|
|
def get_env(self):
|
|
ret_dict = {}
|
|
for k, v in self.__dict__.items():
|
|
if isinstance(v, (list, dict, int, float, str, bool)):
|
|
ret_dict[k] = v
|
|
return ret_dict
|
|
|
|
def set_env(self, we_env):
|
|
for k, v in we_env.items():
|
|
setattr(self, k, v)
|
|
set_random_seed(self.seed)
|
|
|
|
def group_info(self, group):
|
|
group_info = f'group size: {group.size()}\n'
|
|
group_info += f'group rank: {group.rank()}\n'
|
|
group_info += f'group name: {group.name()}\n'
|
|
return group_info
|
|
|
|
@property
|
|
def data_group_world_size(self):
|
|
if self.data_parallel_group is not None:
|
|
return self.data_parallel_group.size()
|
|
return self.world_size
|
|
|
|
@property
|
|
def tensor_group_world_size(self):
|
|
if self.tensor_parallel_group is not None:
|
|
return self.tensor_parallel_group.size()
|
|
return 1
|
|
|
|
@property
|
|
def pipeline_group_world_size(self):
|
|
if self.pipleline_parallel_group is not None:
|
|
return self.pipleline_parallel_group.size()
|
|
return 1
|
|
|
|
def __str__(self):
|
|
environ_str = f'Now running in the distributed environment with world size {self.world_size}\n!'
|
|
environ_str += f'Current pod have {self.device_count} devices!\n'
|
|
environ_str += f'Current task executes on device {self.device_id}!\n'
|
|
environ_str += f"Current task's global rank is {self.rank} \n"
|
|
environ_str += f"Current task's data online is set {self.data_online} \n"
|
|
environ_str += f"Current task's share storage is set {self.share_storage} \n"
|
|
environ_str += f"Current task's global seed is set {self.seed} \n"
|
|
if self.data_parallel_group is not None:
|
|
environ_str += f"Current task's data parallel group: {self.group_info(self.data_parallel_group)} \n"
|
|
if self.pipleline_parallel_group is not None:
|
|
environ_str += f"Current task's pipeline parallel group: {self.group_info(self.pipleline_parallel_group)} \n"
|
|
if self.tensor_parallel_group is not None:
|
|
environ_str += f"Current task's tensor parallel group: {self.group_info(self.tensor_parallel_group)} \n"
|
|
environ_str += f"Current task's backend is set {self.backend} \n"
|
|
return environ_str
|
|
|
|
def __del__(self):
|
|
if we.is_distributed:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
we = Workenv()
|