# run: CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 test.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
torch.distributed.init_process_group(backend="nccl")
 
batch_size = 4
data_size = 8
 
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
 
class RandomDataset(Dataset):
        def __init__(self, length, local_rank):
            self.len = length
            self.data = torch.stack([torch.ones(1), torch.ones(1)*2,torch.ones(1)*3,torch.ones(1)*4,torch.ones(1)*5,torch.ones(1)*6,torch.ones(1)*7,torch.ones(1)*8]).to('cuda')
            self.local_rank = local_rank
        def __getitem__(self, index):
            return self.data[index]
        def __len__(self):
            return self.len
 
dataset = RandomDataset(data_size, local_rank)
sampler = DistributedSampler(dataset)
 
rand_loader = DataLoader(dataset=dataset,batch_size=batch_size,sampler=sampler)
epoch = 0
while epoch < 2:
    sampler.set_epoch(epoch)
    if epoch == 0:
        for data in rand_loader:
                print(data)
    epoch+=1

'''
输出: GPU间的数据不重叠
tensor([[5.],
        [8.],
        [3.],
        [2.]], device='cuda:0')
tensor([[1.],
        [4.],
        [6.],
        [7.]], device='cuda:1')
'''

def _p2p_ops(*,
             tensor_send_prev: Optional[torch.Tensor],
             tensor_recv_prev: Optional[torch.Tensor],
             tensor_send_next: Optional[torch.Tensor],
             tensor_recv_next: Optional[torch.Tensor],
             group: torch.distributed.ProcessGroup):
    reqs = []
    rank = get_pipeline_model_parallel_rank()
    # 状态机
    if get_pipeline_model_parallel_rank() % 2 == 0:
        if tensor_send_next is not None:
            # P2P发送
            send_next_req = torch.distributed.isend(
                tensor=tensor_send_next,
                dst=get_pipeline_model_parallel_next_rank(),
                group=group,
            )
            reqs.append(send_next_req)

        # 状态机
        if tensor_recv_prev is not None:
            # ......

    else:
        if tensor_recv_prev is not None:
            recv_prev_req = torch.distributed.irecv(
                tensor=tensor_recv_prev,
                src=get_pipeline_model_parallel_prev_rank(),
                group=group,
            )
            reqs.append(recv_prev_req)

        # ....
    return reqs

# 相当于p2p_func中构建了状态机的执行计划, 然后会在_communicate中wait()等待通知并做下一步
def _communicate():
    # ...
    # Send tensors in both the forward and backward directions as appropriate.
    if config.use_ring_exchange_p2p:
        def _ring_exchange_wrapper(**kwargs):
            torch.distributed.ring_exchange(**kwargs)
            return []
        p2p_func = _ring_exchange_wrapper
    elif config.batch_p2p_comm:
        assert wait_on_reqs
        p2p_func = _batched_p2p_ops
    else:
        p2p_func = _p2p_ops

    reqs = p2p_func(tensor_send_prev=tensor_send_prev,
                    tensor_recv_prev=tensor_recv_prev,
                    tensor_send_next=tensor_send_next,
                    tensor_recv_next=tensor_recv_next,
                    group=get_pipeline_model_parallel_group())

    if wait_on_reqs and len(reqs) > 0:
        for req in reqs:
            req.wait()
        reqs = None

class ParallelAttention(MegatronModule):
    def __init__(...):
        # ...
        if attention_type == AttnType.self_attn:
            self.query_key_value = tensor_parallel.ColumnParallelLinear(
                config.hidden_size,
                3 * projection_size,
                config=config,
                init_method=config.init_method,
                bias=args.add_bias_linear,
                gather_output=False)

    def forward():
        # ...
        # =====================
        # Query, Key, and Value
        # =====================

        if self.attention_type == AttnType.self_attn:
            # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
            mixed_x_layer, _ = self.query_key_value(hidden_states)

data_parallel_size: int = world_size // (tensor_model_parallel_size *
                                         pipeline_model_parallel_size)

num_tensor_model_parallel_groups: int  = world_size // tensor_model_parallel_size
num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size

# Build the data-parallel groups.
global _DATA_PARALLEL_GROUP
global _DATA_PARALLEL_GLOBAL_RANKS
assert _DATA_PARALLEL_GROUP is None, 'data parallel group is already initialized'
all_data_parallel_group_ranks = []
for i in range(pipeline_model_parallel_size):
    start_rank = i * num_pipeline_model_parallel_groups
    end_rank = (i + 1) * num_pipeline_model_parallel_groups
    for j in range(tensor_model_parallel_size):
        ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
        all_data_parallel_group_ranks.append(list(ranks))
        group = torch.distributed.new_group(ranks)
        if rank in ranks:
            _DATA_PARALLEL_GROUP = group
            _DATA_PARALLEL_GLOBAL_RANKS = ranks

# Build the model-parallel groups.
global _MODEL_PARALLEL_GROUP
assert _MODEL_PARALLEL_GROUP is None, 'model parallel group is already initialized'
for i in range(data_parallel_size):
    ranks = [data_parallel_group_ranks[i]
             for data_parallel_group_ranks in all_data_parallel_group_ranks]
    group = torch.distributed.new_group(ranks)
    if rank in ranks:
        _MODEL_PARALLEL_GROUP = group

# Build the tensor model-parallel groups.
global _TENSOR_MODEL_PARALLEL_GROUP
assert _TENSOR_MODEL_PARALLEL_GROUP is None, \
    'tensor model parallel group is already initialized'
for i in range(num_tensor_model_parallel_groups):
    ranks = range(i * tensor_model_parallel_size,
                  (i + 1) * tensor_model_parallel_size)
    group = torch.distributed.new_group(ranks)
    if rank in ranks:
        _TENSOR_MODEL_PARALLEL_GROUP = group

# Build the pipeline model-parallel groups and embedding groups
# (和Build the model-parallel groups等价)
# (first and last rank in each pipeline model-parallel group).
for i in range(num_pipeline_model_parallel_groups):
    ranks = range(i, world_size, num_pipeline_model_parallel_groups)
    group = torch.distributed.new_group(ranks)
    if rank in ranks:
        _PIPELINE_MODEL_PARALLEL_GROUP = group
        _PIPELINE_GLOBAL_RANKS = ranks
    # Setup embedding group (to exchange gradients between
    # first and last stages).
    if len(ranks) > 1:
        embedding_ranks = [ranks[0], ranks[-1]]
        position_embedding_ranks = [ranks[0]]
        if pipeline_model_parallel_split_rank is not None:
            if ranks[pipeline_model_parallel_split_rank] not in embedding_ranks:
                embedding_ranks = [ranks[0],
                                   ranks[pipeline_model_parallel_split_rank],
                                   ranks[-1]]
            if ranks[pipeline_model_parallel_split_rank] not in position_embedding_ranks:
                position_embedding_ranks = [ranks[0],
                                   ranks[pipeline_model_parallel_split_rank]]
    else:
        embedding_ranks = ranks
        position_embedding_ranks = ranks

    group = torch.distributed.new_group(embedding_ranks)