集合通信操作详解

# Sequence Parallel: 进入 Attention/MLP 前需要完整序列
class _GatherFromSequenceParallelRegion(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input_, group, ...):
        # 前向: AllGather — 每个 TP rank 持有 S/tp_size，拼出完整 S
        return _gather_along_first_dim(input_, group, ...)

    @staticmethod
    def backward(ctx, grad_output):
        # 反向: ReduceScatter — 完整梯度归约后切回 S/tp_size
        return _reduce_scatter_along_first_dim(grad_output, ctx.group, ...)

# Sequence Parallel: 离开 Attention/MLP 后需要切回分片
class _ReduceScatterToSequenceParallelRegion(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input_, group, ...):
        # 前向: ReduceScatter — 各 rank 的部分输出归约后切片
        return _reduce_scatter_along_first_dim(input_, group, ...)

    @staticmethod
    def backward(ctx, grad_output):
        # 反向: AllGather — 梯度分片拼回完整梯度
        return _gather_along_first_dim(grad_output, ctx.group, ...)

def backward(ctx, grad_output):
    input, weight = ctx.saved_tensors
    tp_group = ctx.tp_group

    # ① 计算 input 梯度（完整大小）
    grad_input = grad_output.matmul(weight)            # [S, H]

    if ctx.sequence_parallel:
        # ② 异步 ReduceScatter：把 [S, H] 归约+切分为 [S/tp, H]
        sub_grad_input = torch.empty(dim_size, ...)
        handle = dist_reduce_scatter_func(
            sub_grad_input, grad_input,
            group=tp_group, async_op=True              # ← 异步！不阻塞 GPU
        )
        # 依赖 CUDA_DEVICE_MAX_CONNECTIONS=1 确保 RS 先于 wgrad 入队

    # ③ GPU 继续计算 weight 梯度（计算密集，耗时长）
    # RS 在网卡上传输的同时，GPU 在做这个矩阵乘法
    if ctx.gradient_accumulation_fusion:
        fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
            total_input, grad_output, weight.main_grad  # 融合 wgrad
        )
    else:
        grad_weight = grad_output.t().matmul(total_input)

    # ④ RS 在 ③ 执行期间大概率已完成，此时 wait 几乎不阻塞
    if handle is not None:
        handle.wait()

    return sub_grad_input, grad_weight, ...            # RS 后的分片梯度

# 每个参数注册反向 post-hook
# 当该参数的梯度计算完成后，hook 被自动调用
def _make_backward_post_hook(param):
    def hook(*unused):
        # 梯度累加到 param.main_grad
        param.main_grad.add_(param.grad.data)
        # 标记该参数的梯度已就绪
        bucket_group.register_grad_ready(param)
    return hook

# register_grad_ready 内部：当一个 bucket 内所有参数梯度都就绪时
# 自动触发异步 ReduceScatter
def register_grad_ready(self, param):
    self.params_with_grad.add(param)
    if len(self.params_with_grad) == self.num_params_with_grad:
        self.start_grad_sync()       # 发起异步 RS，async_op=True

# 每个 module 注册前向 pre-hook
# 在该 module 的 forward() 执行之前自动调用
def _make_forward_pre_hook(bucket_group):
    def hook(*unused):
        # ① 等待当前层的 AllGather 完成（参数就绪）
        bucket_group.finish_param_sync()
        # ② 立即发起下一层的异步 AllGather（预取）
        next_bg = bucket_group.next_param_gather_bucket_group
        if next_bg is not None:
            next_bg.start_param_sync()   # async_op=True
    return hook

# 高层 API                           → _communicate() 的参数
# ─────────────────────────────────────────────────────────────────
# recv_forward()                     → recv_prev=True
#   返回 tensor_recv_prev（来自上一 stage 的激活值）

# send_forward(output)               → tensor_send_next=output
#   把前向输出发给下一 stage

# recv_backward()                    → recv_next=True
#   返回 tensor_recv_next（来自下一 stage 的梯度）

# send_backward(grad)                → tensor_send_prev=grad
#   把梯度发给上一 stage

# send_forward_recv_backward(output) → tensor_send_next=output, recv_next=True
#   同时：发激活值给 next + 收梯度从 next（1F1B steady 阶段用）

# send_backward_recv_forward(grad)   → tensor_send_prev=grad, recv_prev=True
#   同时：发梯度给 prev + 收激活值从 prev（1F1B steady 阶段用）

def _communicate(self, *, tensor_send_next, tensor_send_prev,
                 recv_prev, recv_next, tensor_shape, wait_on_reqs=True):
    # ① 如果序列长度可变，先通信 shape
    if config.variable_seq_lengths:
        recv_prev_shape, recv_next_shape = self._communicate_shapes(...)

    # ② 为要接收的 tensor 分配空缓冲区
    if recv_prev:
        tensor_recv_prev = torch.empty(recv_prev_shape, ...)
    if recv_next:
        tensor_recv_next = torch.empty(recv_next_shape, ...)

    # ③ 选择底层实现：三选一
    if config.use_ring_exchange_p2p:     p2p_func = ring_exchange
    elif config.batch_p2p_comm:          p2p_func = _batched_p2p_ops
    else:                                p2p_func = _p2p_ops

    # ④ 调用底层，四个参数全部透传
    p2p_func(tensor_send_prev=..., tensor_recv_prev=...,
             tensor_send_next=..., tensor_recv_next=...,
             group=pp_group, prev_pipeline_rank=..., next_pipeline_rank=...)

    return tensor_recv_prev, tensor_recv_next, reqs

def _batched_p2p_ops(*, tensor_send_prev, tensor_recv_prev,
                     tensor_send_next, tensor_recv_next, group, ...):
    ops = []
    # 把所有需要的 send/recv 操作收集到列表
    if tensor_send_prev is not None:
        ops.append(P2POp(isend, tensor_send_prev, prev_rank, group))
    if tensor_recv_prev is not None:
        ops.append(P2POp(irecv, tensor_recv_prev, prev_rank, group))
    if tensor_send_next is not None:
        ops.append(P2POp(isend, tensor_send_next, next_rank, group))
    if tensor_recv_next is not None:
        ops.append(P2POp(irecv, tensor_recv_next, next_rank, group))

    # 一次性提交，NCCL 内部排序，保证不死锁
    reqs = torch.distributed.batch_isend_irecv(ops)
    return reqs

def _p2p_ops(*, tensor_send_prev, tensor_recv_prev,
             tensor_send_next, tensor_recv_next, group, ...):

    if group.rank() % 2 == 0:    # 偶数 rank：先 send 后 recv
        isend(tensor_send_next, dst=next_rank)    # ① 先发给 next
        irecv(tensor_recv_prev, src=prev_rank)    # ② 再收从 prev
        isend(tensor_send_prev, dst=prev_rank)    # ③ 发给 prev
        irecv(tensor_recv_next, src=next_rank)    # ④ 收从 next

    else:                        # 奇数 rank：先 recv 后 send
        irecv(tensor_recv_prev, src=prev_rank)    # ① 先收从 prev
        isend(tensor_send_next, dst=next_rank)    # ② 再发给 next
        irecv(tensor_recv_next, src=next_rank)    # ③ 收从 next
        isend(tensor_send_prev, dst=prev_rank)    # ④ 发给 prev

even_send_odd_recv_group = group
if group.size() == 2 and get_backend(group) != 'ucc':
    # PP 只有 2 个 stage 时，用全局 WORLD group 做一个方向的通信
    # 这样两个方向的 P2P 可以真正并行，不被同一 group 串行化
    even_recv_odd_send_group = torch.distributed.group.WORLD
else:
    even_recv_odd_send_group = group

# 偶数 rank send 用 even_send_odd_recv_group
# 偶数 rank recv 用 even_recv_odd_send_group（WORLD）
# → 两个通信走不同的 group，可以并行

# Steady 阶段：每次循环做 1 个前向 + 1 个反向
for i in range(num_microbatches_remaining):
    output_tensor = forward_step(input_tensor, ...)

    # 发前向输出（send_next）+ 收反向梯度（recv_next）
    output_tensor_grad = p2p_communicator.send_forward_recv_backward(
        output_tensor, tensor_shapes
    )
    # 底层: _communicate(tensor_send_next=output, recv_next=True)

    input_tensor_grad = backward_step(input_tensor, output_tensor_grad, ...)

    # 发反向梯度（send_prev）+ 收下个前向输入（recv_prev）
    input_tensor = p2p_communicator.send_backward_recv_forward(
        input_tensor_grad, tensor_shapes
    )
    # 底层: _communicate(tensor_send_prev=grad, recv_prev=True)