Tensor Parallelism 实现详解

# f: 前向 Identity，反向 AllReduce
class _CopyToModelParallelRegion(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input_):
        # 前向：直接返回输入，不做任何通信
        # 每个 TP rank 拿到相同的输入 X
        return input_

    @staticmethod
    def backward(ctx, grad_output):
        # 反向：AllReduce 梯度
        # 因为前向时每个 rank 用相同的 X 算了不同的 Y_i
        # 反向时 dL/dX 需要把所有 rank 的梯度贡献求和
        return _reduce(grad_output)

# g: 前向 AllReduce，反向 Identity
class _ReduceFromModelParallelRegion(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input_):
        # 前向：AllReduce 部分结果
        # 每个 TP rank 有 Y_i（部分和），需要 ΣY_i 得到完整输出
        return _reduce(input_)

    @staticmethod
    def backward(ctx, grad_output):
        # 反向：直接返回，不通信
        # 因为前向 AllReduce 后每个 rank 有相同的 Y
        # 下游对 Y 计算的梯度在每个 rank 上也相同
        # 不需要再聚合
        return grad_output

def forward(self, input_):
    # ① 如果开启 SP：输入是 [s/tp, b, h]，需要 AllGather 拼回 [s, b, h]
    if self.sequence_parallel:
        input_parallel = gather_from_sequence_parallel_region(input_, ...)
    else:
        # ② 普通 TP：f 算子 — forward=Identity，直接透传
        input_parallel = copy_to_tensor_model_parallel_region(input_)

    # ③ 矩阵乘法：每个 GPU 用自己的 W_i 分片计算
    # output_parallel: [s, b, h_out/tp]
    output_parallel = linear_with_grad_accumulation(
        input_parallel, self.weight, self.bias, ...
    )

    # ④ 是否收集完整输出
    if self.gather_output:
        # AllGather 拼回完整输出 [s, b, h_out]
        output = gather_from_tensor_model_parallel_region(output_parallel)
    else:
        # 保持切分，直接传给下游 RowParallel
        output = output_parallel

    return output

def forward(self, input_):
    # ① 处理输入
    if self.input_is_parallel:
        # 输入已切分（来自 ColumnParallel 的输出），直接用
        input_parallel = input_
    else:
        # 输入是完整的，Scatter 切分到各 GPU
        input_parallel = scatter_to_tensor_model_parallel_region(input_)

    # ② 矩阵乘法：每个 GPU 用自己的 W_i 分片计算部分和
    # output_parallel: [s, b, h_out]（每个 GPU 得到部分和）
    output_parallel = linear_with_grad_accumulation(
        input_parallel, self.weight, self.bias, ...
    )

    # ③ 如果开启 SP：ReduceScatter（归约+切分到 [s/tp, b, h]）
    if self.sequence_parallel:
        output_ = reduce_scatter_to_sequence_parallel_region(output_parallel, ...)
    else:
        # ④ 普通 TP：g 算子 — AllReduce 部分和
        output_ = reduce_from_tensor_model_parallel_region(output_parallel)

    # ⑤ 加 bias（如果有）
    output = output_ + self.bias if self.bias is not None else output_

    return output

# ColumnParallelLinear.forward 中的 SP 分支
if self.sequence_parallel:
    # SP：输入是 [s/tp, b, h]，先 AllGather 拼回完整序列
    # 用的是 _GatherFromSequenceParallelRegion
    #   forward:  AllGather [s/tp, b, h] → [s, b, h]
    #   backward: ReduceScatter [s, b, h] → [s/tp, b, h]
    input_parallel = gather_from_sequence_parallel_region(
        input_, tensor_parallel_output_grad=...
    )
else:
    # 普通 TP：f 算子 — Identity
    input_parallel = copy_to_tensor_model_parallel_region(input_)

# RowParallelLinear.forward 中的 SP 分支
if self.sequence_parallel:
    # SP：AllReduce 拆成 ReduceScatter
    # 用的是 _ReduceScatterToSequenceParallelRegion
    #   forward:  ReduceScatter [s, b, h] → [s/tp, b, h]
    #   backward: AllGather [s/tp, b, h] → [s, b, h]
    output_ = reduce_scatter_to_sequence_parallel_region(output_parallel, ...)
else:
    # 普通 TP：g 算子 — AllReduce
    output_ = reduce_from_tensor_model_parallel_region(output_parallel)

def forward(self, input_):
    # ① 确定本 GPU 负责的词表范围
    # vocab_start_index = tp_rank * vocab_per_partition
    # vocab_end_index   = vocab_start_index + vocab_per_partition

    # ② 把超出范围的 token ID 映射到 0（避免越界）
    input_mask = (input_ < self.vocab_start_index) | \
                 (input_ >= self.vocab_end_index)
    masked_input = input_.clone() - self.vocab_start_index
    masked_input[input_mask] = 0

    # ③ 查表：用局部偏移后的 ID 查本 GPU 的 embedding 分片
    output_parallel = self.weight[masked_input]

    # ④ 把不属于自己的 token 的输出置零
    output_parallel[input_mask, :] = 0.0

    # ⑤ AllReduce：所有 GPU 的输出求和 → 正确结果
    output = reduce_from_tensor_model_parallel_region(output_parallel)
    return output

def forward(ctx, vocab_parallel_logits, target, label_smoothing):
    # vocab_parallel_logits: [s, b, V/tp] — 每个 GPU 只有部分词表的 logits

    # ① 数值稳定：求全局 max
    logits_max = vocab_parallel_logits.max(dim=-1).values    # 本地 max
    torch.distributed.all_reduce(logits_max,
        op=torch.distributed.ReduceOp.MAX, group=tp_group)  # 全局 max
    vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(-1)

    # ② 计算分母：本地 exp 求和，AllReduce 得到全局 sum(exp)
    exp_logits = vocab_parallel_logits.exp()
    sum_exp = exp_logits.sum(dim=-1)                         # 本地 sum(exp)
    torch.distributed.all_reduce(sum_exp, group=tp_group)    # 全局 sum(exp)

    # ③ 计算分子：只有持有目标 token 的 GPU 贡献 logit 值
    target_mask = (target >= vocab_start) & (target < vocab_end)
    masked_target = target.clone() - vocab_start
    # 查表获取 target 对应的 logit（不持有的 GPU 得到 0）
    predicted_logits = vocab_parallel_logits[..., masked_target] * target_mask
    torch.distributed.all_reduce(predicted_logits, group=tp_group)

    # ④ loss = log(sum_exp) - predicted_logits
    loss = torch.log(sum_exp) - predicted_logits
    return loss