PyTorch在Linux上的并行计算实践

Linux 上 PyTorch 并行计算实践指南

一并行范式与适用场景

数据并行 DataParallel DP：单机多卡，封装模型后在多卡间切分 batch，使用简单但存在单进程瓶颈与显存不均衡问题，适合快速原型验证。
分布式数据并行 DDP DistributedDataParallel：每个 GPU 一个进程，基于NCCL的 All-Reduce 梯度同步，性能与可扩展性更好，是单机与多机训练的主流方案。
模型并行 Model Parallel：将模型的不同层/模块分布到不同设备，解决单卡显存不足，但需显式切分与通信设计。
流水线并行 Pipeline Parallel：将模型按层切分为多个阶段，多卡流水执行，适合超大模型与跨节点扩展。
数据加载并行：使用 DataLoader(num_workers > 0) 的多进程/多线程预取，掩盖 I/O 与预处理开销，常与以上并行方式叠加。
单机多卡场景下，DDP 默认基于 NCCL 通信，通常无需额外 MPI 依赖。

二环境准备与快速检查

驱动与工具链：确认 NVIDIA 驱动与 CUDA 可用，执行 nvidia-smi 查看驱动与 GPU 状态；按需安装 CUDA Toolkit 并设置环境变量（如 PATH、LD_LIBRARY_PATH、CUDA_HOME）。
PyTorch 安装：使用 Conda 安装与 CUDA 版本匹配的 PyTorch，例如：
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia。
多进程与通信：单机多卡优先选择 DDP + NCCL；多机训练需保证节点间网络互通与一致的软件栈。

三实践范例

范例一单机多卡 DDP 最小骨架（torchrun）

# train_ddp.py
import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision.models import resnet18
import torchvision.transforms as T
from torchvision.datasets import CIFAR10

def main():
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl", init_method="env://")

    model = resnet18(num_classes=10).to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    transform = T.Compose([T.ToTensor(), T.Normalize((0.4914,0.4822,0.4465),(0.2023,0.1994,0.2010))])
    train_set = CIFAR10(root="./data", train=True, download=True, transform=transform)
    sampler = DistributedSampler(train_set)
    loader = DataLoader(train_set, batch_size=128, sampler=sampler, num_workers=4, pin_memory=True)

    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    criterion = torch.nn.CrossEntropyLoss()

    model.train()
    for epoch in range(3):
        sampler.set_epoch(epoch)  # 保证每轮shuffle不同
        for x, y in loader:
            x, y = x.to(local_rank, non_blocking=True), y.to(local_rank, non_blocking=True)
            optimizer.zero_grad()
            out = model(x)
            loss = criterion(out, y)
            loss.backward()
            optimizer.step()
        if local_rank == 0:
            print(f"Epoch {epoch} done.")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动方式（单机 4 卡）：
torchrun --nproc_per_node=4 train_ddp.py
要点：使用 DistributedSampler、设置 set_epoch、pin_memory=True、non_blocking 传输、主进程（rank 0）打印与保存。

范例二多机多卡启动（两节点示例）
- 节点 0（IP：192.168.1.1，开放端口 12345）：
  torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr=192.168.1.1 --master_port=12345 train_ddp.py
- 节点 1：
  torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr=192.168.1.1 --master_port=12345 train_ddp.py 说明：--nproc_per_node 为每节点 GPU 数，--nnodes 为节点总数，--node_rank 为节点编号，--master_addr/port 为主节点地址与端口。

范例三数据加载与 CPU 并行

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __len__(self): return 10000
    def __getitem__(self, idx):
        import torch
        return torch.randn(3, 32, 32), torch.randint(0, 10, (1,)).item()

loader = DataLoader(
    MyDataset(),
    batch_size=256,
    num_workers=8,        # 多进程预取
    pin_memory=True,      # 加速 Host->Device 拷贝
    persistent_workers=True  # 长驻 worker 减少重建开销
)

提示：I/O 与预处理密集时增大 num_workers；注意系统 CPU 核心/内存 与磁盘带宽，避免资源争用。

四性能优化与常见坑

并行策略选择
- 单机多卡优先 DDP + NCCL；快速验证可用 DP，但长期训练与扩展性以 DDP 为佳。
- 超大模型采用 模型并行/流水线并行 与 ZeRO 等技术组合。
数据管道
- 使用 DistributedSampler、每轮调用 set_epoch；num_workers 设为 CPU 核心数 的约数，开启 pin_memory 与 non_blocking 传输。
同步与日志
- 多进程仅 rank 0 打印与保存；跨进程指标用 all_reduce 聚合（如 loss/accuracy）。
- 关键阶段使用 torch.cuda.synchronize() 避免计时误差。
批大小与学习率
- DDP 下全局有效 batch 为 per_gpu_batch * world_size，学习率通常按 world_size 线性或按规则缩放。
通信与拓扑
- 优先 NCCL；多机训练确保 高速互联（如 InfiniBand） 与一致的 NCCL/驱动版本。
常见坑
- DP 单进程瓶颈与显存不均；DDP 忘记 set_epoch 导致每轮样本重复；未用 DistributedSampler 导致数据重复或遗漏；多进程重复保存/打印；环境变量 CUDA_VISIBLE_DEVICES 与启动方式不一致。

最新问答

相关标签