在Linux上構建PyTorch分布式訓練環境需要以下幾個步驟:
首先,確保你已經安裝了PyTorch。你可以使用pip或conda來安裝。以下是使用pip安裝的示例:
pip install torch torchvision torchaudio
如果你需要GPU支持,可以安裝相應的CUDA版本:
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
NCCL(NVIDIA Collective Communications Library)是用于多GPU和多節點通信的庫。你可以通過以下命令安裝NCCL:
# 下載并安裝NCCL
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update && sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker
# 下載NCCL庫
wget https://nvidia.github.io/nvidia-docker/$distribution/libcudnn8_8.x.x-1+cuda11.3_amd64.deb
wget https://nvidia.github.io/nvidia-docker/$distribution/libcudnn8-dev_8.x.x-1+cuda11.3_amd64.deb
wget https://nvidia.github.io/nvidia-docker/$distribution/libnccl2_2.x.x-1+cuda11.3_amd64.deb
wget https://nvidia.github.io/nvidia-docker/$distribution/libnccl-dev_2.x.x-1+cuda11.3_amd64.deb
sudo dpkg -i libcudnn8_8.x.x-1+cuda11.3_amd64.deb
sudo dpkg -i libcudnn8-dev_8.x.x-1+cuda11.3_amd64.deb
sudo dpkg -i libnccl2_2.x.x-1+cuda11.3_amd64.deb
sudo dpkg -i libnccl-dev_2.x.x-1+cuda11.3_amd64.deb
為了在分布式訓練中方便地啟動和管理多個節點,你需要配置SSH無密碼登錄。以下是基本步驟:
# 在主節點上生成SSH密鑰
ssh-keygen -t rsa
# 將公鑰復制到所有工作節點
ssh-copy-id user@worker1
ssh-copy-id user@worker2
# 重復以上步驟,直到所有工作節點都配置完畢
使用torch.distributed.launch或accelerate庫來啟動分布式訓練。以下是使用torch.distributed.launch的示例:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=NUM_NODES --node_rank=NODE_RANK --master_addr=MASTER_IP --master_port=12345 your_training_script.py
參數說明:
--nproc_per_node: 每個節點上的GPU數量。--nnodes: 節點總數。--node_rank: 當前節點的排名(從0開始)。--master_addr: 主節點的IP地址。--master_port: 主節點的端口號。accelerate庫(可選)accelerate庫可以簡化分布式訓練的配置和啟動過程。首先安裝accelerate:
pip install accelerate
然后使用accelerate啟動分布式訓練:
accelerate launch your_training_script.py
確保所有節點都能正常通信并且訓練腳本能夠正確運行。你可以通過以下命令檢查節點間的通信:
nccl-tests -b 8 -e 256M -f 2 -g <NUM_GPUS_YOU_HAVE>
通過以上步驟,你應該能夠在Linux上成功構建PyTorch分布式訓練環境。