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 与启动方式不一致。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！