How to operate distributed training of PyTorch on CentOS

PyTorch distributed training on CentOS system requires following the following steps:

1. PyTorch installation: The premise is that Python and pip are installed in CentOS system. Depending on your CUDA version, get the appropriate installation command from the PyTorch official website. For CPU-only training, you can use the following command:

    pip install torch torchvision torchaudio

   If you need GPU support, make sure that the corresponding version of CUDA and cuDNN are installed and use the corresponding PyTorch version to install.

2. Distributed environment configuration: Distributed training usually requires multiple machines or single-machine multiple GPUs. All nodes participating in training must be able to network access to each other and correctly configure environment variables such as MASTER_ADDR (master node IP address) and MASTER_PORT (any available port number).

3. Distributed training script writing: Use PyTorch's torch.distributed package to write distributed training scripts. torch.nn.parallel.DistributedDataParallel is used to wrap your model, while torch.distributed.launch or accelerate libraries are used to start distributed training.

   Here is an example of a simplified distributed training script:

    import torch
   import torch.nn as nn
   import torch.optim as optim
   from torch.nn.parallel import DistributedDataParallel as DDP
   import torch.distributed as dist
   
   def train(rank, world_size):
       dist.init_process_group(backend='nccl', init_method='env://') # Initialize the process group, use the nccl backend model = ... # Your model definition model.cuda(rank) # Move the model to the specified GPU
   
       ddp_model = DDP(model, device_ids=[rank]) # Use DDP to wrap the model criteria = nn.CrossEntropyLoss().cuda(rank) # Loss function optimizer = optim.Adam(ddp_model.parameters(), lr=0.001) # Optimizer dataset = ... # Your dataset sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
       loader = torch.utils.data.DataLoader(dataset, batch_size=..., sampler=sampler)
   
       for epoch in range(...):
           sampler.set_epoch(epoch) # For each epoch resampling, target in loader:
               data, target = data.cuda(rank), target.cuda(rank)
               optimizer.zero_grad()
               output = ddp_model(data)
               loss = criteria(output, target)
               loss.backward()
               optimizer.step()
   
       dist.destroy_process_group() # Destroy process group 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()
       train(args.rank, args.world_size)

4. Distributed training startup: Use the torch.distributed.launch tool to start distributed training. For example, run on two GPUs:

    python -m torch.distributed.launch --nproc_per_node=2 your_training_script.py

   In the case of multiple nodes, ensure that each node runs the corresponding process and that nodes can access each other.

5. Monitoring and debugging: Distributed training may encounter network communication or synchronization problems. Use nccl-tests to test whether the communication between GPUs is normal. Detailed logging is essential for debugging.

Please note that the above steps provide a basic framework that may need to be adjusted according to specific needs and environment in actual applications. It is recommended to refer to the detailed instructions of the official PyTorch documentation on distributed training.

The above is the detailed content of How to operate distributed training of PyTorch on CentOS. For more information, please follow other related articles on the PHP Chinese website!