windows下使用pytorch進行單機多卡分散式訓練

西西嘛呦發表於2023-04-02
WindowsPyTorch分散式

現在有四張卡，但是部署在windows10系統上，想嘗試下在windows上使用單機多卡進行分散式訓練，網上找了一圈硬是沒找到相關的文章。以下是踩坑過程。

首先，pytorch的版本必須是大於1.7，這裡使用的環境是：

pytorch==1.12+cu11.6
四張4090顯示卡
python==3.7.6

使用nn.DataParallel進行分散式訓練

這一種方式較為簡單：
首先我們要定義好使用的GPU的編號，GPU按順序依次為0,1,2,3。gpu_ids可以透過命令列的形式傳入：

gpu_ids = args.gpu_ids.split(',')
gpu_ids = [int(i) for i in gpu_ids]
torch.cuda.set_device('cuda:{}'.format(gpu_ids[0]))

建立模型後用nn.DataParallel進行處理，

 model.cuda()
 r_model = nn.DataParallel(model, device_ids=gpu_ids, output_device=gpu_ids[0])

對，沒錯，只需要這麼兩步就行了。需要注意的是儲存模型後進行載入時，需要先用nn.DataParallel進行處理，再載入權重，不然引數名沒對齊會報錯。

checkpoint = torch.load(checkpoint_path)
model.cuda()
r_model = nn.DataParallel(model, device_ids=gpu_ids, output_device=gpu_ids[0])
r_model.load_state_dict(checkpoint['state_dict'])

如果不使用分散式載入模型，你需要對權重進行對映：

new_start_dict = {}
for k, v in checkpoint['state_dict'].items():
    new_start_dict["module." + k] = v
model.load_state_dict(new_start_dict)

使用Distributed進行分散式訓練

首先了解一下概念：
node：主機數，單機多卡就一個主機，也就是1。
rank：當前程式的序號，用於程式之間的通訊，rank=0的主機為master節點。
local_rank：當前程式對應的GPU編號。
world_size：總的程式數。
在windows中，我們需要在py檔案裡面使用：

import os
os.environ["CUDA_VISIBLE_DEVICES]='0,1,3'

來指定使用的顯示卡。
假設現在我們使用上面的三張顯示卡，執行時顯示卡會重新按照0-N進行編號，有：

[38664] rank = 1, world_size = 3, n = 1, device_ids = [1]
[76032] rank = 0, world_size = 3, n = 1, device_ids = [0]
[23208] rank = 2, world_size = 3, n = 1, device_ids = [2]

也就是程式0使用第1張顯示卡，進行1使用第2張顯示卡，程式2使用第三張顯示卡。
有了上述的基本知識，再看看具體的實現。

使用torch.distributed.launch啟動

使用torch.distributed.launch啟動時，我們必須要在args裡面新增一個local_rank引數，也就是：
parser.add_argument("--local_rank", type=int, default=0)
1、初始化：

import torch.distributed as dist

env_dict = {
        key: os.environ[key]
        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")
}
current_work_dir = os.getcwd()
    init_method = f"file:///{os.path.join(current_work_dir, 'ddp_example')}"
dist.init_process_group(backend="gloo", init_method=init_method, rank=int(env_dict["RANK"]),
                                world_size=int(env_dict["WORLD_SIZE"]))

這裡需要重點注意，這種啟動方式在環境變數中會存在RANK和WORLD_SIZE，我們可以拿來用。backend必須指定為gloo，init_method必須是file:///，而且每次執行完一次，下一次再執行前都必須刪除生成的ddp_example，不然會一直卡住。
2、構建模型並封裝
local_rank會自己繫結值，不再是我們--local_rank指定的。

 model.cuda(args.local_rank)
 r_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=device_ids)

3、構建資料集載入器並封裝

  train_dataset = dataset(file_path='data/{}/{}'.format(args.data_name, train_file))
  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  train_loader = DataLoader(train_dataset, batch_size=args.train_batch_size,
                              collate_fn=collate.collate_fn, num_workers=4, sampler=train_sampler)

4、計算損失函式
我們把每一個GPU上的loss進行匯聚後計算。

def loss_reduce(self, loss):
        rt = loss.clone()
        dist.all_reduce(rt, op=dist.ReduceOp.SUM)
        rt /= self.args.local_world_size
        return rt

loss = self.criterion(outputs, labels)
torch.distributed.barrier()
loss = self.loss_reduce(loss)

注意列印相關資訊和儲存模型的時候我們通常只需要在local_rank=0時列印。同時，在需要將張量轉換到GPU上時，我們需要指定使用的GPU，透過local_rank指定就行，即data.cuda(args.local_rank)，保證資料在對應的GPU上進行處理。
5、啟動
在windows下需要把換行符去掉，且只變為一行。

python -m torch.distributed.launch \
--nnode=1 \
--node_rank=0 \
--nproc_per_node=3 \
main_distributed.py \
--local_world_size=3 \
--bert_dir="../model_hub/chinese-bert-wwm-ext/" \
--data_dir="./data/cnews/" \
--data_name="cnews" \
--log_dir="./logs/" \
--output_dir="./checkpoints/" \
--num_tags=10 \
--seed=123 \
--max_seq_len=512 \
--lr=3e-5 \
--train_batch_size=64 \
--train_epochs=1 \
--eval_batch_size=64 \
--do_train \
--do_predict \
--do_test

nproc_per_node、local_world_size和GPU的數目保持一致。

使用torch.multiprocessing啟動

使用torch.multiprocessing啟動和使用torch.distributed.launch啟動大體上是差不多的，有一些地方需要注意。

mp.spawn(main_worker, nprocs=args.nprocs, args=(args,))

main_worker是我們的主執行函式，dist.init_process_group要放在這裡面，而且第一個引數必須為local_rank。即main_worker(local_rank, args)
nprocs是程式數，也就是使用的GPU數目。
args按順序傳入main_worker真正使用的引數。
其餘的就差不多。
啟動指令：

python main_mp_distributed.py \
--local_world_size=4 \
--bert_dir="../model_hub/chinese-bert-wwm-ext/" \
--data_dir="./data/cnews/" \
--data_name="cnews" \
--log_dir="./logs/" \
--output_dir="./checkpoints/" \
--num_tags=10 \
--seed=123 \
--max_seq_len=512 \
--lr=3e-5 \
--train_batch_size=64 \
--train_epochs=1 \
--eval_batch_size=64 \
--do_train \
--do_predict \
--do_test

最後需要說明的，假設我們設定的batch_size=64，那麼實際上的batch_size = int(batch_size / GPU數目)。
附上完整的基於bert的中文文字分類單機多卡訓練程式碼：https://github.com/taishan1994/pytorch_bert_chinese_text_classification

參考

https://github.com/tczhangzhi/pytorch-distributed
https://murphypei.github.io/blog/2020/09/pytorch-distributed
https://pytorch.org/docs/master/distributed.html?highlight=all_gather#torch.distributed.all_gather
https://github.com/lesliejackson/PyTorch-Distributed-Training
https://github.com/pytorch/examples/blob/ddp-tutorial-code/distributed/ddp/example.py
996黃金一代：[原創][深度][PyTorch] DDP系列第一篇：入門教程
「新生手冊」：PyTorch分散式訓練 - 知乎 (zhihu.com)

相關文章