PyTorch在Ubuntu上的分布式训练指南

以下是在Ubuntu上进行PyTorch分布式训练的指南：

环境准备

• 确保所有节点安装相同版本的PyTorch、CUDA和NCCL，可使用pip install torch torchvision torchaudio安装PyTorch。
• 确保节点间可通过SSH无密码登录，配置网络使节点能相互通信，防火墙允许相关端口通信。

编写训练脚本

• 初始化进程组：在脚本中添加dist.init_process_group，指定backend（如nccl）、init_method（如tcp://主节点IP:端口）、world_size和rank。
• 模型并行：用DistributedDataParallel包装模型，如model = DDP(model, device_ids=[rank])。
• 数据加载：使用DistributedSampler包装数据集，如train_sampler = DistributedSampler(train_dataset)。

启动分布式训练

• 使用torch.distributed.launch，命令为python -m torch.distributed.launch --nproc_per_node=GPU数量 --nnodes=节点数 --node_rank=当前节点排名 --master_addr=主节点IP --master_port=端口号 脚本.py。
• 也可使用accelerate库，先安装pip install accelerate，再运行accelerate launch 脚本.py。