2. 张量并行基础

任务目标

实现线性层（Y = X @ W^T + b）的按行切分张量并行版本。将权重矩阵 W 按行切分为两块，分别放置在两个设备上，每个设备独立完成部分矩阵乘法后，通过 AllReduce 求和得到与单卡完全一致的输出。

核心原理

按行切分时，W 被切分为 [W1; W2]（上下两块），输入 X 需完整复制到每个设备。每个设备计算 Y_i = X @ W_i^T，得到的是输出特征维度的部分和（各设备输出维度不同），因此需要 AllReduce 将各设备的输出在特征维度上拼接（此处为求和后拼接，等价于 AllReduce Sum）。

环境依赖

• Python 3.8+
• PyTorch 2.0+（含 torch.distributed）
• 单机多卡（或 CPU 模拟多进程）

完整实现

import os
import torch
import torch.distributed as dist
import torch.nn.functional as F

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group('gloo', rank=rank, world_size=world_size)

def row_parallel_linear(rank, world_size, X, W, b):
    """按行切分线性层: Y = X @ W^T + b"""
    # 按行切分权重 (输出维度切分)
    d_out = W.size(0)
    assert d_out % world_size == 0, "输出维度需能被卡数整除"
    local_d_out = d_out // world_size
    W_local = W[rank * local_d_out : (rank + 1) * local_d_out].clone()
    b_local = b[rank * local_d_out : (rank + 1) * local_d_out].clone()

    # 每卡持有完整输入 X，计算本地部分输出
    Y_local = X @ W_local.T + b_local  # [batch, local_d_out]

    # AllReduce Sum: 各卡输出维度不同，需先拼接再求和
    # 收集所有卡的部分输出
    gathered = [torch.zeros_like(Y_local) for _ in range(world_size)]
    dist.all_gather(gathered, Y_local)
    Y_full = torch.cat(gathered, dim=1)  # [batch, d_out]
    return Y_full

def verify(rank, world_size):
    setup(rank, world_size)
    torch.manual_seed(42)

    batch, d_in, d_out = 4, 8, 6
    X = torch.randn(batch, d_in)
    W = torch.randn(d_out, d_in)
    b = torch.randn(d_out)

    # 单卡参考结果
    Y_ref = X @ W.T + b

    # 张量并行结果
    Y_parallel = row_parallel_linear(rank, world_size, X, W, b)

    # 验证一致性
    if rank == 0:
        assert torch.allclose(Y_parallel, Y_ref, atol=1e-5), "并行结果与单卡不一致!"
        print(f"验证通过: 最大误差 = {(Y_parallel - Y_ref).abs().max().item():.2e}")
    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = 2
    rank = int(os.environ.get('RANK', 0))
    verify(rank, world_size)

运行方式

# 单机双卡（或 CPU 多进程模拟）
torchrun --nproc_per_node=2 row_parallel.py

验证方法与判定标准

• 程序正常退出，无报错。
• rank 0 输出 验证通过: 最大误差 = x.xxe-xx，且误差小于 1e-5。
• 若将 world_size 改为 1 运行，结果与单卡完全一致（自洽性检查）。

关键说明

• 按行切分要求 d_out 能被 world_size 整除，否则需 padding 或改用其他切分策略。
• AllReduce 在此实现中通过 all_gather + cat 完成，实际工程中可用 AllGather 集合通信原语直接实现。
• 输入 X 需完整复制到每卡（本实现中由外部传入同一份数据，实际分布式场景需广播或每卡持有相同输入）。