|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.distributed as dist |
|
|
from .comm.pg_utils import ProcessGroupManager |
|
|
from .comm.comm import set_sp_comm_group, split_sequence, gather_sequence, all_to_all_comm |
|
|
from .comm.operation import gather_forward_split_backward |
|
|
|
|
|
class SequenceParallelManager: |
|
|
_SP_GROUP = None |
|
|
_SP_SIZE = 0 |
|
|
|
|
|
@staticmethod |
|
|
def sp_on(): |
|
|
return SequenceParallelManager._SP_GROUP is not None |
|
|
|
|
|
@staticmethod |
|
|
def init_sp(sp_size): |
|
|
if SequenceParallelManager._SP_GROUP is not None: |
|
|
print("WARN: sequence parallel group is already initialized") |
|
|
return |
|
|
|
|
|
if sp_size <= 1: |
|
|
print(f"WARN: sequence parallel size must > 1 but got {sp_size}") |
|
|
return |
|
|
|
|
|
world_size = dist.get_world_size() |
|
|
assert world_size % sp_size == 0, f"world_size {world_size} must be divisible by sp_size({sp_size})" |
|
|
SequenceParallelManager._SP_SIZE = sp_size |
|
|
|
|
|
pm = ProcessGroupManager( |
|
|
world_size // sp_size, |
|
|
sp_size, |
|
|
dp_axis=0, |
|
|
sp_axis=1, |
|
|
) |
|
|
pm_group = pm.sp_group |
|
|
set_sp_comm_group(pm_group) |
|
|
SequenceParallelManager._SP_GROUP = pm_group |
|
|
return |
|
|
|
|
|
@staticmethod |
|
|
def get_sp_group(): |
|
|
return SequenceParallelManager._SP_GROUP |
|
|
|
|
|
@staticmethod |
|
|
def get_sp_size(): |
|
|
return SequenceParallelManager._SP_SIZE |
|
|
|
|
|
@staticmethod |
|
|
def get_sp_group_nums(): |
|
|
|
|
|
if SequenceParallelManager.sp_on(): |
|
|
world_size = torch.distributed.get_world_size() |
|
|
return world_size // SequenceParallelManager._SP_SIZE |
|
|
else: |
|
|
return 0 |
|
|
|
|
|
@staticmethod |
|
|
def get_sp_rank(): |
|
|
if SequenceParallelManager.sp_on(): |
|
|
global_rank = torch.distributed.get_rank() |
|
|
sp_rank = global_rank % SequenceParallelManager._SP_SIZE |
|
|
return sp_rank |
|
|
else: |
|
|
return 0 |
|
|
|
|
|
def get_sp_group_rank(): |
|
|
if SequenceParallelManager.sp_on(): |
|
|
global_rank = torch.distributed.get_rank() |
|
|
sp_group_rank = global_rank // SequenceParallelManager._SP_SIZE |
|
|
return sp_group_rank |
|
|
else: |
|
|
return 0 |
|
|
|
|
|
def sp_split_sequence_by_dim(seq, seqlen_dim=1) -> torch.Tensor: |
|
|
""" |
|
|
split the raw sequence by seqlen_dim |
|
|
""" |
|
|
return split_sequence(seq, SequenceParallelManager.get_sp_group(), seqlen_dim, 'down') |
|
|
|
|
|
def sp_gather_sequence_by_dim(seq, seqlen_dim=1) -> torch.Tensor: |
|
|
""" |
|
|
gather seqlen_dim to recover raw sequence |
|
|
""" |
|
|
return gather_sequence(seq, SequenceParallelManager.get_sp_group(), seqlen_dim, 'up') |
|
|
|
|
|
def sp_all_to_all(ts, scatter_dim, gather_dim): |
|
|
""" |
|
|
reorder the tensor's dimension, like [raw_seq_len/sp_size, hidden_dim] to [raw_seq_len, hidden_dim/sp_size] |
|
|
|
|
|
scatter_dim: the dimension to split the tensor |
|
|
gather_dim: the dimension to concatenate |
|
|
""" |
|
|
|
|
|
return all_to_all_comm(ts, SequenceParallelManager.get_sp_group(), scatter_dim, gather_dim) |
|
|
|
|
|
|
