如何在Ubuntu上配置PyTorch的网络环境

在 Ubuntu 上配置 PyTorch 的网络环境

一 目标与准备

• 明确目标：在 Ubuntu 上搭建可复现的 PyTorch 开发/训练环境，优先使用 NVIDIA GPU 加速，并准备好后续的多机分布式训练网络条件。
• 基础检查与准备：
  • 确认 NVIDIA 显卡可被识别：lspci | grep -i nvidia
  • 更新系统并安装常用工具：sudo apt update && sudo apt upgrade -y && sudo apt install -y build-essential cmake git wget
  • 建议安装 Anaconda 管理 Python 与依赖：bash Anaconda3-*.sh，安装后 source ~/anaconda3/bin/activate 并确认命令行前缀出现 (base)。

二 安装驱动与 CUDA 工具链

• 显卡驱动与 CUDA 版本选择
  • 查看驱动支持的最高 CUDA：nvidia-smi（右上显示 Supported/Runtime CUDA）。
  • 原则：驱动向后兼容，优先选择较新的稳定驱动；随后安装的 CUDA Toolkit 与 cuDNN 需与驱动和 PyTorch 版本匹配。
• 安装方式（二选一）
  • 推荐：仅安装驱动（run 文件或 Ubuntu 附加驱动），随后按需安装 CUDA Toolkit 与 cuDNN（便于独立升级与回滚）。
  • 备选：安装 CUDA Toolkit 时勾选驱动（版本较单一，灵活性差）。
• 安装后验证
  • 驱动：nvidia-smi 正常显示 GPU 型号、驱动版本、CUDA Runtime。
  • 工具链：nvcc -V 显示 CUDA 编译器版本；如需 cuDNN，按官方指引放置头文件与库并验证链接。

三 创建 Conda 环境并安装 PyTorch

• 创建隔离环境（示例 Python 3.8）
  • conda create -n pytorch_env python=3.8
  • conda activate pytorch_env
• 安装 PyTorch（以官网命令为准，以下为常见示例）
  • CPU 版：conda install pytorch torchvision torchaudio cpuonly -c pytorch
  • GPU 版（示例 CUDA 11.1）：conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch
  • 若网络不稳定，重试或使用国内镜像源（如清华/中科大 conda 镜像）。
• 验证安装
  • python -c "import torch; print('PyTorch:', torch.__version__)"
  • GPU 可用性：python -c "import torch; print('CUDA available:', torch.cuda.is_available())"。

四 单机多卡与多机分布式网络配置

• 单机多卡（DataParallel/DistributedDataParallel）
  • 使用 NCCL 作为后端通常性能更佳；确保 NCCL 与驱动、CUDA 版本匹配。
  • 启动示例（DDP）：python -m torch.distributed.launch --nproc_per_node=2 --use_env train.py
• 多机多卡（以 TCP 初始化为例）
  • 主节点 IP 与端口需可达（如 192.168.1.10:12355），所有节点防火墙放行该端口。
  • 指定通信网卡（常见为 eth0/enpXs0 等）：export GLOO_SOCKET_IFNAME=enp5s0（如使用 Gloo 后端）；如使用 NCCL，可设置 NCCL_SOCKET_IFNAME。
  • 启动命令（各节点）：python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=192.168.1.10 --master_port=12355 train.py
  • 常见问题：端口被占用、网卡名错误、跨网段路由不通、防火墙未放行。

五 常见问题与快速排查

• nvidia-smi 正常但 torch.cuda.is_available() 为 False
  • 常见于 PyTorch 与 CUDA/cuDNN 版本不匹配；请按 PyTorch 官网选择与驱动匹配的 cudatoolkit 版本重新安装。
• 多机通信失败
  • 核对 master_addr/master_port、各节点 rank 与 world_size；使用 netstat -tulpen | grep 12355 检查端口监听；必要时设置 GLOO_SOCKET_IFNAME 或 NCCL_SOCKET_IFNAME 指定网卡。
• 包下载慢或中断
  • 使用 conda 镜像源 或在稳定网络环境下重试安装命令。