上一篇文章: Pytorch 分散式訓練 DistributedDataParallel — 概念篇 有介紹分散式訓練的概念，本文將要來進行 Pytorch DistributedDataParallel 實作。
首先來介紹一下分散式的基礎概念
- group: 指進程組，默認為一組。
- backend: 指進程使用的通訊後端，Pytorch 支援 mpi、gloo、nccl，若是使用 Nvidia GPU 推薦使用 nccl。
關於後端的詳細資訊可由官方文檔 DISTRIBUTED COMMUNICATION PACKAGE — TORCH.DISTRIBUTED 查看。
- world_size: 指進程組中的進程數量。
若使用單台機器多 GPU ，world_size 表示使用的 GPU 數量
若使用多台機器多 GPU ，world_size 表示使用的機器數量
- rank: 指當前進程的序號，用於進程間的通訊，rank=0 表示為 master。
若使用單台機器多 GPU ，rank表示當前正在使用的某個 GPU
若使用多台機器多 GPU ，rank表示當前正在使用的某台機器
- local_rank: 指每台機器上的進程序號。
📚 rank vs local_rank
假設目前使用兩台機器，每台機器有 3 個 GPU，每個 GPU 使用一個進程 (DDP 的最佳使用方式是每一個 GPU 使用一個進程)。
機器一的 rank 序號為 0, 1, 2；機器二的 rank 序號為 3, 4, 5
機器一的 local_rank 序號為 0, 1, 2；機器二的 local_rank 序號為 0, 1, 2
接下來就來開始實作啦~ 先 import 需要的 library，我的 pytorch 版本為 1.9.0、torchvision 版本為 0.10.0。
import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler設置 local_rank argparse 參數
在啟動分散式訓練時，需要在命令行使用 torch.distributed.launch 啟動器，該啟動器會將當前進程的序號 (若每個 GPU 使用一個進程，也是指 GPU 序號) 通過 local_rank 參數傳給 python 檔。
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()初始化 DDP
設置進程通訊的後端為 nccl
可參考官方參數設置: https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group
dist.init_process_group(backend='nccl')同步所有進程
dist.barrier()獲取當前進程組內的所有進程數
world_size = dist.get_world_size()數據平行化
將數據集劃分給每一個進程，避免進程間的數據重複
train_sampler = DistributedSampler(train_dataset)
valid_sampler = DistributedSampler(valid_dataset)- DistributedSampler 分配方式
假設使用 2 個 GPU 且數據集分為 0~4 份，首先將數據集進行 shuffle 後分為兩份，不足的部分由第一份數據補 (第 5 份)，最後再以間隔方式分配數據至不同的 GPU 中。
接著再將 sampler 傳入 DataLoader
DataLoader 官方參數設置: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)
模型
假設模型已經 load，直接使用 .to(device)，將模型分配至 CUDA
device = torch.device("cuda", args.local_rank)
model = model.to(device)使用 SyncBN 優化各自進程中數據集的 BN 計算
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)使用 DistributedDataParallel 將模型包裝起來
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)set_epoch
在每一次迭代前，通過 set_epoch 來設置一個 random seed，這樣每次運行才會將每個 epoch 的數據都 shuffle，並且讓每個 GPU 拿到的數據不同。
train_sampler.set_epoch(epoch)
valid_sampler.set_epoch(epoch)啟動分散式訓練
這部分有兩種寫法: torch.distributed.launch、torchrun
1. torch.distributed.launch
*參數說明
nproc_per_node: 每台機器有多少進程
nnodes: 總共使用多少台機器
node_rank: 當前進程位於哪台機器
master_addr: master 進程的網路地址
master_port: master 進程的端口，要先確認該端口是否被其他進程佔用
可參考官方設置: https://pytorch.org/docs/stable/distributed.html#launch-utility
- 運行一台機器多 GPU
CUDA_VISIBLE_DEVICES 參數為選擇要使用哪一台 GPU，0,1 表示使用第 0、1 序號的 GPU
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py- 運行多台機器多 GPU
假設 IP 為 192.168.1.1、free port 為 1234
# 第一台機器
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=”192.168.1.1" --master_port=1234 ddp_example.py# 第二台機器
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=”192.168.1.1" --master_port=1234 ddp_example.py
- 運行一台機器多 GPU，但執行兩個不同實驗
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py# 使用不同 master_port 避免衝突
$ CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 --master_port 9999 ddp_example.py
2. torchrun
Pytorch 提供了新的指令替代 torch.distributed.launch，參數的使用方式可參考官方文件: https://pytorch.org/docs/stable/elastic/run.html#launcher-api
- 運行一台機器多 GPU
$ CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc_per_node=2 ddp_example.py- 運行多台機器多 GPU
參數 nnodes 為使用機器序號的區間、rdzv_id 為唯一的 job id、rdzv_endpoint 為 backend 所執行的 endpoint，格式為 <host>:<port> ，替代原本的 master_addr 及 master_port。
$ torchrun --nnodes=1:4 --nproc_per_node=$NUM_TRAINERS --rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR ddp_example.py完整的 Code 可參考我的 Github:
執行以下指令進行訓練
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py顯示結果如下
最後，總結 Pytorch 分散式訓練的基本流程:
- 設置 local_rank 參數並使用 init_process_group 進行初始化
- 使用 DistributedSampler 劃分數據集
- 將模型分配至 CUDA 並設置 SyncBN 及 DistributedDataParallel
- 將 train、vaild Sampler 設置 set_epoch
- 在 sever 上執行指令開始訓練
