mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
2.8 KiB
74 lines
2.8 KiB
3 years ago
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
|
||
|
|
from colossalai.context.parallel_mode import ParallelMode
|
||
|
|
from colossalai.core import global_context as gpc
|
||
|
|
from colossalai.utils import get_current_device
|
||
|
|
|
||
|
|
|
||
|
|
def send_tensor_meta(tensor, need_meta=True, down_group=None):
|
||
|
|
"""Sends tensor meta information before sending a specific tensor.
|
||
|
|
Since the recipient must know the shape of the tensor in p2p communications,
|
||
|
|
meta information of the tensor should be sent before communications. This function
|
||
|
|
synchronizes with :func:`recv_tensor_meta`.
|
||
|
|
|
||
|
|
:param tensor: Tensor to be sent
|
||
|
|
:param need_meta: If False, meta information won't be sent
|
||
|
|
:param down_group: Communication group including the next member in pipeline parallel group
|
||
|
|
:type tensor: Tensor
|
||
|
|
:type need_meta: bool, optional
|
||
|
|
:type down_group: ProcessGroup, optional
|
||
|
|
:return: False
|
||
|
|
:rtype: bool
|
||
|
|
"""
|
||
|
|
if need_meta:
|
||
|
|
rank = gpc.get_global_rank()
|
||
|
|
|
||
|
|
if down_group is None:
|
||
|
|
down_group = gpc.get_group(ParallelMode.PIPELINE_NEXT)
|
||
|
|
|
||
|
|
tensor_kwargs = {'dtype': torch.long, 'device': get_current_device()}
|
||
|
|
|
||
|
|
send_shape = torch.tensor(tensor.size(), **tensor_kwargs)
|
||
|
|
send_ndims = torch.tensor(len(tensor.size()), **tensor_kwargs)
|
||
|
|
|
||
|
|
dist.broadcast(send_ndims, src=rank, group=down_group)
|
||
|
|
dist.broadcast(send_shape, src=rank, group=down_group)
|
||
|
|
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def recv_tensor_meta(tensor_shape, prev_rank=None, up_group=None):
|
||
|
|
"""Recieves tensor meta information before recieving a specific tensor.
|
||
|
|
Since the recipient must know the shape of the tensor in p2p communications,
|
||
|
|
meta information of the tensor should be recieved before communications. This function
|
||
|
|
synchronizes with :func:`send_tensor_meta`.
|
||
|
|
|
||
|
|
:param tensor_shape: The shape of the tensor to be recieved
|
||
|
|
:param prev_rank: The rank of the source of the tensor
|
||
|
|
:param up_group: Communication group including the previous member in pipeline parallel group
|
||
|
|
:type tensor_shape: torch.Size
|
||
|
|
:type prev_rank: int, optional
|
||
|
|
:type up_group: ProcessGroup, optional
|
||
|
|
:return: The shape of the tensor to be recieved
|
||
|
|
:rtype: torch.Size
|
||
|
|
"""
|
||
|
|
if tensor_shape is None:
|
||
|
|
if prev_rank is None:
|
||
|
|
prev_rank = gpc.get_prev_global_rank(
|
||
|
|
ParallelMode.PIPELINE)
|
||
|
|
if up_group is None:
|
||
|
|
up_group = gpc.get_group(ParallelMode.PIPELINE_PREV)
|
||
|
|
|
||
|
|
tensor_kwargs = {'dtype': torch.long, 'device': get_current_device()}
|
||
|
|
|
||
|
|
recv_ndims = torch.empty((), **tensor_kwargs)
|
||
|
|
dist.broadcast(recv_ndims, src=prev_rank, group=up_group)
|
||
|
|
|
||
|
|
recv_shape = torch.empty(recv_ndims, **tensor_kwargs)
|
||
|
|
dist.broadcast(recv_shape, src=prev_rank, group=up_group)
|
||
|
|
|
||
|
|
tensor_shape = torch.Size(recv_shape)
|
||
|
|
|
||
|
|
return tensor_shape
|