#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from typing import List, Tuple, Union, Any
import pickle
import io
import torch
import torch.distributed as dist
from torch.distributed import distributed_c10d as c10d
from torch.distributed import ProcessGroupNCCL
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
TensorShape = Union[torch.Size, List[int], Tuple[int]]
_pg_manager = {}
_unpickler = pickle.Unpickler
def init_process_group():
"""intialise process group by dist.new_group in the adjacent stages
Args:
None
Returns:
None
"""
world_size = gpc.get_world_size(ParallelMode.PIPELINE)
for i in range(world_size - 1):
_pg_manager[(i, i + 1)] = dist.new_group([i, i + 1])
def _acquire_pair_group_handle(first_rank: int, second_rank: int) -> ProcessGroupNCCL:
"""get the group handle of two given ranks
Args:
first_rank (int): first rank in the pair
second_rank (int): second rank in the pair
Returns:
:class:`ProcessGroupNCCL`: the handle of the group consisting of the given two ranks
"""
if len(_pg_manager) == 0:
init_process_group()
if first_rank > second_rank:
first_rank, second_rank = second_rank, first_rank
pair_key = (first_rank, second_rank)
return _pg_manager[pair_key]
def _cuda_safe_tensor_to_object(tensor: torch.Tensor, tensor_size: torch.Size) -> object:
"""transform tensor to object with unpickle.
Info of the device in bytes stream will be modified into current device before unpickling
Args:
tensor (:class:`torch.tensor`): tensor to be unpickled
tensor_size (:class:`torch.Size`): Size of the real info in bytes
Returns:
Any: object after unpickled
"""
buf = tensor.numpy().tobytes()[:tensor_size]
if b'cuda' in buf:
buf_array = bytearray(buf)
device_index = torch.cuda.current_device()
buf_array[buf_array.find(b'cuda') + 5] = 48 + device_index
buf = bytes(buf_array)
io_bytes = io.BytesIO(buf)
byte_pickler = _unpickler(io_bytes)
unpickle = byte_pickler.load()
return unpickle
def _broadcast_object_list(object_list: List[Any], src: int, dst: int, device=None):
"""This is a modified version of the broadcast_object_list in torch.distribution
The only difference is that object will be move to correct device after unpickled.
If local_rank = src, then object list will be sent to rank src. Otherwise, object list will
be updated with data sent from rank src.
Args:
object_list (List[Any]): list of object to broadcast
src (int): source rank to broadcast
dst (int): dst rank to broadcast
device (:class:`torch.device`): device to do broadcast. current device in default
"""
group = _acquire_pair_group_handle(src, dst)
if c10d._rank_not_in_group(group):
c10d._warn_not_in_group("broadcast_object_list")
return
local_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
# Serialize object_list elements to tensors on src rank.
if local_rank == src:
tensor_list, size_list = zip(*[c10d._object_to_tensor(obj) for obj in object_list])
object_sizes_tensor = torch.cat(size_list)
else:
object_sizes_tensor = torch.empty(len(object_list), dtype=torch.long)
is_nccl_backend = c10d._check_for_nccl_backend(group)
current_device = None
if device is not None:
if is_nccl_backend and device.type != "cuda":
raise ValueError("device type must be cuda for nccl backend")
current_device = device
else:
current_device = torch.device("cpu")
if is_nccl_backend:
current_device = torch.device("cuda", torch.cuda.current_device())
if is_nccl_backend:
object_sizes_tensor = object_sizes_tensor.to(current_device)
# Broadcast object sizes
c10d.broadcast(object_sizes_tensor, src=src, group=group, async_op=False)
# Concatenate and broadcast serialized object tensors
if local_rank == src:
object_tensor = torch.cat(tensor_list)
else:
object_tensor = torch.empty( # type: ignore[call-overload]
torch.sum(object_sizes_tensor).item(), # type: ignore[arg-type]
dtype=torch.uint8,
)
if is_nccl_backend:
object_tensor = object_tensor.to(current_device)
c10d.broadcast(object_tensor, src=src, group=group, async_op=False)
# Deserialize objects using their stored sizes.
offset = 0
if local_rank != src:
for i, obj_size in enumerate(object_sizes_tensor):
obj_view = object_tensor[offset:offset + obj_size]
obj_view = obj_view.type(torch.uint8)
if obj_view.device != torch.device("cpu"):
obj_view = obj_view.cpu()
offset += obj_size
# unpickle
unpickle_object = _cuda_safe_tensor_to_object(obj_view, obj_size)
# unconsistence in device
if isinstance(unpickle_object,
torch.Tensor) and unpickle_object.device.index != torch.cuda.current_device():
unpickle_object = unpickle_object.cuda()
object_list[i] = unpickle_object
def _send_object(object: Any, dst: int) -> None:
"""send anything to dst rank
Args:
object (Any): object needed to be sent
dst (int): rank of the destination
Returns:
None
"""
local_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
# handler = _acquire_pair_group_handle(local_rank, dst)
# transform to list if not
if isinstance(object, torch.Tensor):
object = [object]
# broadcast length first
# TODO : more elegant ? P.S. reduce a _broadcast_object_list
_broadcast_object_list([len(object)], local_rank, dst)
# then broadcast safely
_broadcast_object_list(object, local_rank, dst)
def _recv_object(src: int) -> Any:
"""recv anything from src
Args:
src (int): source rank of data. local rank will receive data from src rank.
Returns:
Any: Object received from src.
"""
local_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
# handler = _acquire_pair_group_handle(local_rank, src)
# recv length first
length = [0]
_broadcast_object_list(length, src, local_rank)
# then create recv buff from length[0] and broadcast
object = [None] * length[0]
_broadcast_object_list(object, src, local_rank)
if length[0] == 1:
object = object[0]
return object
def recv_forward(prev_rank: int = None) -> Any:
"""Copy the forward output from the previous stage in pipeline as the input tensor of this stage.
Args:
input_tensor_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): The shape of the tensor to be received.
prev_rank (int, optional): The rank of the source of the tensor.
Returns:
Any: The input tensor or input tensor list.
"""
if gpc.is_pipeline_first_stage():
input_tensor = None
else:
if prev_rank is None:
prev_rank = gpc.get_prev_global_rank(ParallelMode.PIPELINE)
input_tensor = _recv_object(prev_rank)
return input_tensor
def recv_backward(next_rank: int = None) -> Any:
"""Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage.
Args:
output_grad_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): The shape of the tensor to be received.
next_rank (int, optional): The rank of the source of the tensor.
Returns:
Any: The input gradient tensor or gradident tensor list.
"""
if gpc.is_pipeline_last_stage():
output_tensor_grad = None
else:
if next_rank is None:
next_rank = gpc.get_next_global_rank(ParallelMode.PIPELINE)
output_tensor_grad = _recv_object(next_rank)
return output_tensor_grad
def send_forward(output_object: Any, next_rank: int = None) -> None:
"""Sends the input tensor to the next stage in pipeline.
Args:
output_object Any: Object to be sent.
next_rank (int, optional): The rank of the recipient of the tensor.
"""
if not gpc.is_pipeline_last_stage():
if next_rank is None:
next_rank = gpc.get_next_global_rank(ParallelMode.PIPELINE)
_send_object(output_object, next_rank)
def send_backward(input_object: Any, prev_rank: int = None) -> None:
"""Sends the gradient tensor to the previous stage in pipeline.
Args:
input_tensor_grad (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): Tensor to be sent
prev_rank (int, optional): The rank of the recipient of the tensor
"""
if not gpc.is_pipeline_first_stage():
if prev_rank is None:
prev_rank = gpc.get_prev_global_rank(ParallelMode.PIPELINE)
_send_object(input_object, prev_rank)