3D 并行的实现：TP / PP / DP 各占几行

nanotron 的最大教育价值是"3D 并行实现都很短"。本章把 TP / PP / DP 的核心代码摘出来看。

5.1TP：TensorParallelColumnLinear（~100 行）

看 src/nanotron/parallel/tensor_parallel/nn.py：

class TensorParallelColumnLinear(nn.Linear):
    def __init__(self, in_features, out_features, pg, mode, ...):
        # 切 output 维度
        local_out = out_features // pg.size()
        super().__init__(in_features, local_out, bias=bias)
        self.pg = pg                       # tp process group
        self.mode = mode                    # ALL_REDUCE / REDUCE_SCATTER

    def forward(self, x):
        # 1) input：所有 rank 完整
        if self.mode == "ALL_REDUCE":
            y = F.linear(x, self.weight, self.bias)
            return y                       # output 是 sharded（每 rank 1/N 列）
        else:  # REDUCE_SCATTER
            x_full = all_gather(x, self.pg)
            y = F.linear(x_full, self.weight, self.bias)
            return y

class TensorParallelRowLinear(nn.Linear):
    def __init__(self, in_features, out_features, pg, mode, ...):
        local_in = in_features // pg.size()
        super().__init__(local_in, out_features, bias=bias)
        ...

    def forward(self, x):
        # x 是 input 维切的
        y = F.linear(x, self.weight, self.bias)
        if self.mode == "ALL_REDUCE":
            y = all_reduce(y, self.pg)     # 完整 output
        else:
            y = reduce_scatter(y, self.pg)
        return y

就这么简单。Llama 的 Attention 用 ColumnLinear (q/k/v) + RowLinear (o)，MLP 类似。

5.2PP：1F1B schedule（~200 行）

看 src/nanotron/parallel/pipeline_parallel/engine.py：

class OneForwardOneBackwardPipelineEngine:
    def step(self, model, data_iter, criterion, optimizer):
        num_microbatches = config.num_microbatches
        pp_size = pg.size()

        # warm-up：每个 stage 跑 (pp_size - pp_rank - 1) 个 forward
        warmup_steps = pp_size - pp_rank - 1
        for _ in range(warmup_steps):
            input = recv_from_prev_stage()
            output = forward(model, input)
            send_to_next_stage(output)

        # 1F1B：每跑一个 forward 跑一个 backward
        for _ in range(num_microbatches - warmup_steps):
            input = recv_from_prev_stage()
            output = forward(model, input)
            send_to_next_stage(output)
            grad = recv_grad_from_next_stage()
            input_grad = backward(output, grad)
            send_grad_to_prev_stage(input_grad)

        # cool-down：剩下的 backward
        for _ in range(warmup_steps):
            grad = recv_grad_from_next_stage()
            input_grad = backward(output, grad)
            send_grad_to_prev_stage(input_grad)

实际代码有更多细节（buffer 管理 / loss accumulation），但核心就这个 schedule。

5.3DP：DistributedDataParallel + 梯度 all-reduce

nanotron 没用 PyTorch 自带的 DDP，而是自己实现了一个：

# src/nanotron/parallel/data_parallel/utils.py
def sync_gradients(model, dp_pg):
    """backward 完后调一次，把每张卡的梯度 all-reduce 求平均"""
    for p in model.parameters():
        if p.grad is not None:
            dist.all_reduce(p.grad, op=dist.ReduceOp.AVG, group=dp_pg)

核心就一个 all-reduce 循环。生产代码还有梯度 bucket、async 通信等优化。

5.4三种并行怎么组合

看 src/nanotron/trainer.py 训练循环：

def training_step(self, batch):
    if self.pp_size > 1:
        loss = self.pp_engine.step(self.model, batch, self.criterion, self.optimizer)
    else:
        with self.tp_context():
            output = self.model(batch)
            loss = self.criterion(output, batch.target)
            loss.backward()

    # DP 同步梯度
    if self.dp_size > 1:
        self.sync_gradients()

    self.optimizer.step()
    self.optimizer.zero_grad()

清晰的三层：TP 内部 forward；PP schedule 控制 stage 之间；DP 同步梯度。

5.5TP mode：ALL_REDUCE vs REDUCE_SCATTER

大模型 / 长序列推荐 REDUCE_SCATTER 模式（即 Megatron 的 SP）。

5.6这章你需要带走的

• TP ColumnLinear / RowLinear 各 ~100 行，跟 Megatron 思路一致；
• PP 1F1B 主流程 ~50 行：warmup → 1F1B → cool-down；
• DP 自家实现 = 一个梯度 all-reduce 循环；
• TP mode 选 REDUCE_SCATTER + SP 在长序列省 activation；
• 读 nanotron 的 TP/PP 源码是入门分布式训练的最佳途径。