1. 张量并行基础

任务目标

实现一个线性层（不带bias）的按列切分（Column Parallel）版本，将权重矩阵沿输出维度切分为两块，分别部署到两个GPU上。前向时两个GPU各自计算部分输出后拼接得到完整输出；反向时梯度按对应切片回传。最终通过数值对比验证前向与反向结果与单卡完整线性层一致。

环境与依赖

• Python 3.8+，PyTorch 2.0+（需支持 torch.distributed）
• 两个GPU（或单机多卡），可通过 torchrun 启动

核心实现

import torch
import torch.nn as nn
import torch.distributed as dist

class ColumnParallelLinear(nn.Module):
    """按列切分线性层：每个rank持有 weight[:, local_out_start:local_out_end]"""
    def __init__(self, in_features, out_features, rank, world_size):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.world_size = world_size

        # 每个rank的本地输出维度
        assert out_features % world_size == 0, "out_features必须能被world_size整除"
        self.local_out_features = out_features // world_size

        # 本地权重切片，形状: [in_features, local_out_features]
        self.weight = nn.Parameter(torch.randn(in_features, self.local_out_features) * 0.1)

    def forward(self, x):
        # x: [batch, in_features]
        local_out = torch.matmul(x, self.weight)  # [batch, local_out_features]
        return local_out

def run_worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    batch, in_f, out_f = 4, 8, 6
    torch.manual_seed(42)

    # 完整权重（仅用于验证，实际并行时不存储）
    full_weight = torch.randn(in_f, out_f) * 0.1
    # 按列切分
    local_out_size = out_f // world_size
    local_weight = full_weight[:, rank*local_out_size : (rank+1)*local_out_size].cuda()

    # 构建并行层并同步本地权重
    layer = ColumnParallelLinear(in_f, out_f, rank, world_size).cuda()
    with torch.no_grad():
        layer.weight.copy_(local_weight)

    # 输入（所有rank相同）
    x = torch.randn(batch, in_f).cuda()

    # 前向
    local_out = layer(x)  # [batch, local_out_features]

    # 收集所有rank的local_out并拼接
    gathered = [torch.zeros_like(local_out) for _ in range(world_size)]
    dist.all_gather(gathered, local_out)
    if rank == 0:
        full_out = torch.cat(gathered, dim=1)  # [batch, out_f]
        ref_out = torch.matmul(x.cpu(), full_weight)  # 单卡参考
        assert torch.allclose(full_out.cpu(), ref_out, atol=1e-5), "前向结果不一致"
        print("前向验证通过")

    # 反向：构造随机梯度并验证
    grad_out = torch.randn(batch, out_f).cuda()
    local_grad = grad_out[:, rank*local_out_size : (rank+1)*local_out_size]

    # 并行层反向
    layer.zero_grad()
    local_out.backward(local_grad)
    local_grad_w = layer.weight.grad  # [in_f, local_out_features]

    # 参考反向
    ref_grad_w = torch.matmul(x.cpu().T, grad_out.cpu())  # [in_f, out_f]
    ref_local_grad_w = ref_grad_w[:, rank*local_out_size : (rank+1)*local_out_size]

    assert torch.allclose(local_grad_w.cpu(), ref_local_grad_w, atol=1e-5), "反向梯度不一致"
    if rank == 0:
        print("反向验证通过")

    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = 2
    torch.multiprocessing.spawn(run_worker, args=(world_size,), nprocs=world_size)

运行方式

torchrun --nproc_per_node=2 column_parallel_linear.py

验证方法与判定标准

• 前向验证：rank 0 将两个GPU的局部输出按列拼接后，与单卡完整权重计算结果对比，误差小于 1e-5 即通过。
• 反向验证：每个rank将本地梯度与参考梯度对应切片对比，误差小于 1e-5 即通过。
• 程序正常结束且打印“前向验证通过”和“反向验证通过”即视为测试通过。

关键说明

• 本实现假设 out_features 能被 world_size 整除，且未包含bias以聚焦核心切分逻辑。
• 实际工程中需在 forward 后调用 AllGather 获取完整输出，本案例为验证正确性在rank 0上执行拼接对比。
• 反向传播时，PyTorch自动计算本地权重的梯度，无需额外通信；若需对输入求梯度，则需在反向时执行 AllReduce，本案例未涉及。