PyTorch Linux集群搭建方法是什么

PyTorch Linux集群搭建方法

一 架构与准备

• 硬件与网络：至少准备2台 Linux 服务器，每台含1–8 张 NVIDIA GPU；节点间建议使用千兆以上以太网或 InfiniBand，并尽量使用**共享存储（如 NFS）**统一存放数据与代码，便于多机读取一致性。所有节点需在同一局域网，且时间同步（如使用 NTP）。
• 软件环境：统一 Python 3.8–3.10、安装匹配版本的 PyTorch + CUDA + cuDNN，GPU 训练优先使用 NCCL 2.10+；如需跨平台或 CPU 通信，可备选 Gloo。
• 账号与权限：建议配置SSH 免密登录（主节点可免密登录所有工作节点），便于统一拉起训练进程。
• 驱动与连通性：安装/验证 NVIDIA 驱动 与 nvidia-smi 正常；开放用于初始化通信的 MASTER_PORT（如 29500） 的防火墙端口。

二 环境部署与连通性验证

• 基础环境（示例）：
  • 安装 Miniconda/Anaconda，创建隔离环境：conda create -n ddp python=3.9
  • 安装 PyTorch（示例）：conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
  • 验证：python -c “import torch; print(torch.cuda.is_available())”（应返回 True）
• 集群侧准备：
  • 主节点生成 SSH 密钥并分发公钥到各工作节点，实现免密登录（ssh-keygen; ssh-copy-id user@ip）。
  • 开放通信端口（示例）：sudo ufw allow 29500/tcp；测试端口：telnet 29500。
  • 共享数据：将数据集放在 NFS 挂载目录（如 /mnt/nfs/data），各节点以相同路径访问。

三 训练脚本要点（DDP）

• 初始化进程组：使用 init_process_group 指定后端（GPU 优先 “nccl”，跨平台或 CPU 用 “gloo”），通过 TCP 初始化（init_method=“tcp://MASTER_ADDR:MASTER_PORT”），设置 rank 与 world_size。
• 模型与设备：模型放到对应 GPU（to(rank % torch.cuda.device_count())），用 DistributedDataParallel 包装；DataLoader 使用 DistributedSampler 并每轮设置 sampler.set_epoch(epoch)。
• 启动方式：优先使用 torchrun 统一拉起多进程/多节点任务（推荐），也可使用 mp.spawn 或旧版 torch.distributed.launch。
• 最小示例（要点）：
  • 初始化：dist.init_process_group(backend=“nccl”, init_method=“tcp://MASTER_ADDR:MASTER_PORT”, world_size=WORLD_SIZE, rank=RANK)
  • 模型：model = model.to(rank); ddp_model = DDP(model, device_ids=[rank])
  • 数据：sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank); dataloader = DataLoader(dataset, batch_size=…, sampler=sampler)
  • 训练循环：sampler.set_epoch(epoch); 前向/反向/优化；结束后 dist.destroy_process_group()

四 启动方式与作业调度

• 使用 torchrun（推荐，弹性扩缩）：
  • 主节点（node_rank=0）：torchrun --nnodes=N --nproc_per_node=G --node_rank=0 --master_addr=MASTER_IP --master_port=29500 train.py
  • 工作节点（node_rank=i）：torchrun --nnodes=N --nproc_per_node=G --node_rank=i --master_addr=MASTER_IP --master_port=29500 train.py
  • 说明：N 为节点总数，G 为每节点 GPU 数；环境变量 RANK/LOCAL_RANK/WORLD_SIZE 会被自动注入。
• 使用 Slurm（HPC 集群常用）：
  • 示例脚本要点：
    • 资源申请：#SBATCH -N 2 --ntasks-per-node=8 --gres=gpu:8
    • 环境加载：source ~/miniconda3/etc/profile.d/conda.sh; conda activate ddp; module load mpi/hpcx …
    • 启动命令：srun --mpi=pmix python -u train.py（或结合 torchrun 使用）
• 旧版/备选方式：
  • mp.spawn 在单机多卡或多进程场景可用；
  • torch.distributed.launch 仍可见于旧项目，但官方更推荐 torchrun。

五 常见问题与优化

• 连接失败/超时：检查 MASTER_ADDR/MASTER_PORT、防火墙、节点间 ping/端口连通性；必要时设置 GLOO_SOCKET_IFNAME 指定网卡，或调整 GLOO_TIMEOUT_SECONDS 防止超时。
• 性能优化：优先 NCCL + 万兆/IB；每 GPU 合理设置 batch_size 与 num_workers；启用 梯度累积 维持有效总批量；使用 TensorBoard 与 nvidia-smi 持续观测吞吐与显存。
• 数据与一致性：多机务必使用 共享存储 或确保各节点数据路径一致；分布式 Sampler 必须设置 set_epoch(epoch) 保证洗牌一致性。
• 稳定性：训练前先在单机多卡验证逻辑，再扩展到多机；合理设置 timeout_mins，并定期 checkpoint 以便中断恢复。