1. 首頁 > 生活百科 > >

使用Pytorch進行多卡訓練( 二 )

生活百科


5、初始化子進程組 ， 定義進程間的通信后端（還有GLOO、MPI ， 只有NCCL支持GPU間通信）、子進程rank、子進程數量 。只有當該函數在size個進程中被調用時 ， 各進程才會繼續從這里執行下去 。這個函數統一了各子進程后續代碼的開始時間 。
6、執行子進程代碼 。
7、由于創建子進程會執行本程序，因此主進程的執行需要放在__main__里，防止子進程執行 。
8、開始創建子進程的方式：spawn、fork 。windows默認spawn，linux默認fork 。具體區別請百度 。
9、由于是以NCCL為通信后端的分布式訓練，如果不同進程中相同名稱的張量在同一GPU上，當這個張量進行進程間通信時就會出錯 。為了防止出錯 ， 限制每張卡獨占一個進程，每個進程獨占一張卡 。這里有兩張卡，所以最多只能創建兩個進程 。
10、創建子進程，傳入子進程的初始化方法，及子進程調用該方法的參數 。
11、等待子進程全部運行完畢后再退出主進程 。
輸出結果如下：

使用Pytorch進行多卡訓練

文章插圖
正是各進程保存在不同GPU上的張量的廣播求和(all_reduce)的結果 。
參考： https://pytorch.org/tutorials/intermediate/dist_tuto.html
Pytorch分布式訓練DEMO我們實際上可以根據上面的分布式基礎寫一個分布式訓練，但由于不知道pytorch如何實現GPU間模型梯度的求和 ， 即官方教程中所謂的ring_reduce（沒找到相關API），時間原因，就不再去搜索相關方法了 。這里僅記錄pytorh內部的分布式模型訓練，即利用DDP模塊實現 。Pytorch版本1.12.1 。
import torch,osimport torch.distributed as distimport torch.multiprocessing as mpimport torch.optim as optimfrom torch.nn.parallel import DistributedDataParallel as DDPfrom torch import nndef example(rank, world_size):dist.init_process_group("nccl", rank=rank, world_size=world_size)# ——1——model = nn.Linear(2, 1, False).to(rank)if rank == 0: # ——2——model.load_state_dict(torch.load('model_weight'))# model_stat = torch.load('model_weight', {'cuda:0':'cuda:%d'%rank})#這樣讀取保險一點# model.load_state_dict(model_stat)opt = optim.Adam(model.parameters(), lr=0.0001) # ——3——opt_stat = torch.load('opt_weight', {'cuda:0':'cuda:%d'%rank}) # ——4——opt.load_state_dict(opt_stat) # ——5——ddp_model = DDP(model, device_ids=[rank])# ——6inp = torch.tensor([[1.,2]]).to(rank) # ——7——labels = torch.tensor([[5.]]).to(rank)outp = ddp_model(inp)loss = torch.mean((outp - labels)**2)opt.zero_grad()loss.backward() # ——8——opt.step() # ——9if rank == 0:# ——10——torch.save(model.state_dict(), 'model_weight')torch.save(opt.state_dict(), 'opt_weight')if __name__=="__main__":os.environ["MASTER_ADDR"] = "localhost"# ——11——os.environ["MASTER_PORT"] = "29500"world_size = 2mp.spawn(example, args=(world_size,), nprocs=world_size, join=True) # ——12——以上代碼包含模型在多GPU上讀取權重、進行分布式訓練、保存權重等過程 。細節注釋如下：
1、初始化進程組，由于使用GPU通信，后端應該寫為NCCL 。不過經過實驗，即使錯寫為gloo ， DDP內部也會自動使用NCCL作為通信模塊 。
2、由于后面使用DDP包裹模型進行訓練 ， 其內部會自動將所有rank的模型權重同步為rank 0的權重，因此我們只需在rank 0上讀取模型權重即可 。這是基于Pytorch版本1.12.1，低級版本似乎沒有這個特性，需要在不同rank分別導入權重，則load需要傳入map_location，如下面注釋的兩行代碼所示 。
3、這里創建model的優化器，而不是創建用ddp包裹后的ddp_model的優化器 ， 是為了兼容單GPU訓練，讀取優化器權重更方便 。
4、將優化器權重讀取至該進程占用的GPU 。如果沒有map_location參數，load會將權重讀取到原本保存它時的設備 。
5、優化器獲取權重 。經過實驗，即使權重不在優化器所在的GPU，權重也會遷移過去而不會報錯 。當然load直接讀取到相應GPU會減少數據傳輸 。
6、DDP包裹模型，為模型復制一個副本到相應GPU中 。所有rank的模型副本會與rank 0保持一致 。注意，DDP并不復制模型優化器的副本，因此各進程的優化器需要我們在初始化時保持一致 。權重要么不讀取，要么都讀取 。
7、這里開始模型的訓練 。數據需轉移到相應的GPU設備 。
8、在backward中，所有進程的模型計算梯度后 ， 會進行平均（不是相加） 。也就是說，DDP在backward函數添加了hook ， 所有進程的模型梯度的ring_reduce將在這里執行 。這個可以通過給各進程模型分別輸入不同的數據進行驗證 ， backward后這些模型有相同的梯度，且驗算的確是所有進程梯度的平均 。此外 ， 還可以驗證backward函數會阻斷(block)各進程使用梯度，只有當所有進程都完成backward之后，各進程才能讀取和使用梯度 。這保證了所有進程在梯度上的一致性 。

推薦閱讀