centos

PyTorch在CentOS上的分布式训练如何实现

小樊
49
2025-06-12 16:07:32
栏目: 智能运维

在CentOS上使用PyTorch进行分布式训练,可以按照以下步骤进行:

环境准备

  1. 安装Python和PyTorch

    • 确保CentOS系统上已经安装了Python。
    • 使用pip或conda安装PyTorch。例如:
      pip install torch torchvision torchaudio
      
  2. 安装依赖库

    • 安装必要的系统依赖库,如ncclmpi4py等。
      sudo yum install -y epel-release
      sudo yum install -y libnccl-devel
      pip install mpi4py
      
  3. 配置SSH无密码登录

    • 在所有节点之间配置SSH无密码登录,以便分布式训练时可以无缝通信。

分布式训练设置

  1. 启动分布式训练环境

    • 使用mpiruntorch.distributed.launch来启动分布式训练。
    • 例如,使用torch.distributed.launch
      python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py
      
  2. 编写分布式训练脚本

    • 在训练脚本中,使用torch.distributed.init_process_group来初始化分布式环境。
    • 示例代码:
      import torch
      import torch.nn as nn
      import torch.optim as optim
      from torch.nn.parallel import DistributedDataParallel as DDP
      
      def main(rank, world_size):
          torch.manual_seed(1234)
          torch.cuda.set_device(rank)
      
          # 初始化分布式环境
          torch.distributed.init_process_group(
              backend='nccl',
              init_method='tcp://<master_ip>:<master_port>',
              world_size=world_size,
              rank=rank
          )
      
          # 创建模型并移动到GPU
          model = nn.Linear(10, 10).to(rank)
          ddp_model = DDP(model, device_ids=[rank])
      
          # 创建损失函数和优化器
          criterion = nn.CrossEntropyLoss()
          optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
      
          # 训练循环
          for epoch in range(10):
              optimizer.zero_grad()
              inputs = torch.randn(20, 10).to(rank)
              labels = torch.randint(0, 10, (20,)).to(rank)
              outputs = ddp_model(inputs)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
              print(f'Rank {rank}, Epoch {epoch}, Loss {loss.item()}')
      
      if __name__ == '__main__':
          import argparse
          parser = argparse.ArgumentParser()
          parser.add_argument('--world_size', type=int, default=2)
          parser.add_argument('--rank', type=int, default=0)
          args = parser.parse_args()
      
          main(args.rank, args.world_size)
      
  3. 运行分布式训练

    • 在每个节点上运行上述脚本,并确保world_sizerank参数正确设置。
    • 例如,在两个节点上运行:
      # 节点1
      python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py --world_size=2 --rank=0
      
      # 节点2
      python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py --world_size=2 --rank=1
      

注意事项

通过以上步骤,你可以在CentOS上实现PyTorch的分布式训练。

0
看了该问题的人还看了