|
|
|
|
|
|
|
|
import os |
|
|
|
|
|
import torch.distributed as dist |
|
|
|
|
|
|
|
|
class COMM_INFO: |
|
|
|
|
|
def __init__(self): |
|
|
self.group = None |
|
|
self.sp_size = 1 |
|
|
self.global_rank = 0 |
|
|
self.rank_within_group = 0 |
|
|
self.group_id = 0 |
|
|
|
|
|
|
|
|
nccl_info = COMM_INFO() |
|
|
_SEQUENCE_PARALLEL_STATE = False |
|
|
|
|
|
|
|
|
def initialize_sequence_parallel_state(sequence_parallel_size): |
|
|
global _SEQUENCE_PARALLEL_STATE |
|
|
if sequence_parallel_size > 1: |
|
|
_SEQUENCE_PARALLEL_STATE = True |
|
|
initialize_sequence_parallel_group(sequence_parallel_size) |
|
|
else: |
|
|
nccl_info.sp_size = 1 |
|
|
nccl_info.global_rank = int(os.getenv("RANK", "0")) |
|
|
nccl_info.rank_within_group = 0 |
|
|
nccl_info.group_id = int(os.getenv("RANK", "0")) |
|
|
|
|
|
|
|
|
def set_sequence_parallel_state(state): |
|
|
global _SEQUENCE_PARALLEL_STATE |
|
|
_SEQUENCE_PARALLEL_STATE = state |
|
|
|
|
|
|
|
|
def get_sequence_parallel_state(): |
|
|
return _SEQUENCE_PARALLEL_STATE |
|
|
|
|
|
|
|
|
def initialize_sequence_parallel_group(sequence_parallel_size): |
|
|
"""Initialize the sequence parallel group.""" |
|
|
rank = int(os.getenv("RANK", "0")) |
|
|
world_size = int(os.getenv("WORLD_SIZE", "1")) |
|
|
assert ( |
|
|
world_size % sequence_parallel_size == 0 |
|
|
), "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format( |
|
|
world_size, sequence_parallel_size) |
|
|
nccl_info.sp_size = sequence_parallel_size |
|
|
nccl_info.global_rank = rank |
|
|
num_sequence_parallel_groups: int = world_size // sequence_parallel_size |
|
|
for i in range(num_sequence_parallel_groups): |
|
|
ranks = range(i * sequence_parallel_size, |
|
|
(i + 1) * sequence_parallel_size) |
|
|
group = dist.new_group(ranks) |
|
|
if rank in ranks: |
|
|
nccl_info.group = group |
|
|
nccl_info.rank_within_group = rank - i * sequence_parallel_size |
|
|
nccl_info.group_id = i |
|
|
|
|
|
|
|
|
def destroy_sequence_parallel_group(): |
|
|
"""Destroy the sequence parallel group.""" |
|
|
dist.destroy_process_group() |
|
|
|
