Ubuntu分布式训练环境搭建（手把手教你配置PyTorch多机多卡深度学习训练环境）

一、准备工作

在开始之前，请确保你有以下资源：

• 至少两台运行 Ubuntu 20.04 或更高版本的服务器（或虚拟机）
• 每台机器配备 NVIDIA GPU（建议驱动版本 ≥ 470）
• 稳定的内网连接（建议千兆或万兆网络）
• Python 3.8+ 和 pip 已安装

二、安装基础依赖

首先，在所有节点上执行以下命令安装必要的软件包：

# 更新系统sudo apt update && sudo apt upgrade -y# 安装基础工具sudo apt install -y python3-pip git openssh-server net-tools# 启动 SSH 服务（用于节点间通信）sudo systemctl enable sshsudo systemctl start ssh

三、配置 GPU 驱动与 CUDA

在每台机器上安装 NVIDIA 驱动和 CUDA Toolkit。推荐使用官方 NVIDIA 驱动：

# 添加官方仓库wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.debsudo dpkg -i cuda-keyring_1.0-1_all.debsudo apt update# 安装 CUDA 11.8（可根据需求调整版本）sudo apt install -y cuda-toolkit-11-8# 将 CUDA 加入环境变量echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrcecho 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrcsource ~/.bashrc

验证安装是否成功：

nvidia-smi  # 应显示 GPU 信息nvcc --version  # 应显示 CUDA 编译器版本

四、安装 PyTorch 与分布式依赖

使用 pip 安装支持 GPU 的 PyTorch。确保所有节点安装相同版本：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

此外，安装 NCCL（NVIDIA Collective Communications Library），这是 PyTorch 分布式训练的核心通信库：

sudo apt install -y libnccl2 libnccl-dev

五、配置节点间免密 SSH 登录

为方便管理，建议配置主节点到其他节点的免密登录。假设主节点 IP 为 192.168.1.10，工作节点为 192.168.1.11 和 192.168.1.12。

# 在主节点生成 SSH 密钥（如未生成）ssh-keygen -t rsa -b 4096# 将公钥复制到所有工作节点ssh-copy-id user@192.168.1.11ssh-copy-id user@192.168.1.12# 测试免密登录ssh user@192.168.1.11 'echo "Success!"'

六、编写分布式训练脚本

下面是一个简单的 PyTorch 分布式训练示例（使用 torch.distributed.launch）：

# train_ddp.pyimport torchimport torch.distributed as distimport torch.multiprocessing as mpfrom torch.nn.parallel import DistributedDataParallel as DDPimport osdef setup(rank, world_size):    os.environ['MASTER_ADDR'] = '192.168.1.10'  # 主节点 IP    os.environ['MASTER_PORT'] = '12355'    dist.init_process_group("nccl", rank=rank, world_size=world_size)def cleanup():    dist.destroy_process_group()def main(rank, world_size):    setup(rank, world_size)        # 创建模型并放到对应 GPU    model = torch.nn.Linear(10, 1).to(rank)    ddp_model = DDP(model, device_ids=[rank])        # 简单训练循环    loss_fn = torch.nn.MSELoss()    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.001)        for epoch in range(10):        optimizer.zero_grad()        outputs = ddp_model(torch.randn(20, 10).to(rank))        labels = torch.randn(20, 1).to(rank)        loss = loss_fn(outputs, labels)        loss.backward()        optimizer.step()        if rank == 0:            print(f"Epoch {epoch}, Loss: {loss.item()}")        cleanup()if __name__ == "__main__":    world_size = 2  # 总 GPU 数量（跨节点）    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)

七、启动分布式训练

在主节点上运行以下命令启动训练（假设有两台机器，每台有1个GPU）：

# 在主节点执行python3 -m torch.distributed.run \  --nproc_per_node=1 \  --nnodes=2 \  --node_rank=0 \  --master_addr="192.168.1.10" \  --master_port=12355 \  train_ddp.py# 在工作节点（192.168.1.11）执行python3 -m torch.distributed.run \  --nproc_per_node=1 \  --nnodes=2 \  --node_rank=1 \  --master_addr="192.168.1.10" \  --master_port=12355 \  train_ddp.py

八、常见问题排查

• 防火墙阻塞：确保 12355 端口（或你指定的端口）在所有节点开放。
• 版本不一致：所有节点的 PyTorch、CUDA、NCCL 版本必须完全一致。
• IP 配置错误：检查 MASTER_ADDR 是否为主节点内网 IP。

结语

通过以上步骤，你已经成功搭建了一个基于 Ubuntu 的深度学习环境搭建平台，能够高效运行PyTorch分布式训练任务。随着模型规模不断增长，掌握多机多卡训练技能将成为 AI 工程师的核心竞争力。希望本教程能为你提供清晰、实用的指导！