0. Abstract

使用 PyTorch 進(jìn)行多卡訓(xùn)練, 最簡(jiǎn)單的是 DataParallel, 僅僅添加一兩行代碼就可以使模型在多張 GPU 上并行地計(jì)算. 但它是比較老的方法, 官方推薦使用新的 Distributed Data Parallel, 更加靈活與強(qiáng)大:

1. Distributed Data Parallel (DDP)

從一個(gè)簡(jiǎn)單的非分布式訓(xùn)練任務(wù), 到多機(jī)器多卡訓(xùn)練. 跟著官方教程走, 剛開(kāi)始一切都很順利, 到最后要多機(jī)器的時(shí)候, 就老是報(bào)錯(cuò): MemoryError: std::bad_alloc.

1.1 DDP 概覽

特點(diǎn):

• 多個(gè) batch 的數(shù)據(jù), 同時(shí)分別在多個(gè) GPU 上計(jì)算;
• 需要 DistributedSampler 給各 GPU 分發(fā)數(shù)據(jù) batch, 保證數(shù)據(jù)不重復(fù);
• 模型在各 GPU 上都有一份副本, 分別計(jì)算梯度, 并通過(guò) ring all-reduce 算法整合梯度.

可以理解為: 為每個(gè) GPU 啟動(dòng)一個(gè)進(jìn)程, 這些進(jìn)程執(zhí)行著完全相同的代碼(你的程序), 不同的地方在于:

• 吃進(jìn)了不同的數(shù)據(jù)樣本, 那么計(jì)算得到的 loss 和反向傳播計(jì)算的參數(shù)梯度都不同;
• 各進(jìn)程有自己的編號(hào)(rank), 程序中可根據(jù)編號(hào)執(zhí)行一些不同的操作, 如保存 checkpoint, 日志輸出等操作.

1.2 single_gpu.py

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from datautils import MyTrainDatasetclass Trainer:def __init__(self,model: torch.nn.Module,train_data: DataLoader,optimizer: torch.optim.Optimizer,gpu_id: int,save_every: int,) -> None:self.gpu_id = gpu_idself.model = model.to(gpu_id)self.train_data = train_dataself.optimizer = optimizerself.save_every = save_everydef _run_batch(self, source, targets):output = self.model(source)loss = F.cross_entropy(output, targets)self.optimizer.zero_grad()loss.backward()self.optimizer.step()def _run_epoch(self, epoch):b_sz = len(next(iter(self.train_data))[0])print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")for source, targets in self.train_data:source = source.to(self.gpu_id)targets = targets.to(self.gpu_id)self._run_batch(source, targets)def _save_checkpoint(self, epoch):ckp = self.model.state_dict()PATH = "checkpoint.pt"torch.save(ckp, PATH)print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")def train(self, max_epochs: int):for epoch in range(max_epochs):self._run_epoch(epoch)if epoch % self.save_every == 0:self._save_checkpoint(epoch)def load_train_objs():train_set = MyTrainDataset(2048)  # load your datasetmodel = torch.nn.Linear(20, 1)  # load your modeloptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)return train_set, model, optimizerdef prepare_dataloader(dataset: Dataset, batch_size: int):return DataLoader(dataset,batch_size=batch_size,pin_memory=True,shuffle=True)def main(device, total_epochs, save_every, batch_size):dataset, model, optimizer = load_train_objs()train_data = prepare_dataloader(dataset, batch_size)trainer = Trainer(model, train_data, optimizer, device, save_every)trainer.train(total_epochs)if __name__ == "__main__":import argparseparser = argparse.ArgumentParser(description='simple distributed training job')parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')parser.add_argument('save_every', type=int, help='How often to save a snapshot')parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')args = parser.parse_args()device = 0  # shorthand for cuda:0main(device, args.total_epochs, args.save_every, args.batch_size)

1.3 multi_gpu.py

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from datautils import MyTrainDatasetimport torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import osdef ddp_setup(rank, world_size):"""Args:rank: Unique identifier of each processworld_size: Total number of processes"""# MASTER 表示主節(jié)點(diǎn), 負(fù)責(zé)分配任務(wù), 啟動(dòng)其他進(jìn)程os.environ["MASTER_ADDR"] = "localhost"  # IP address of masteros.environ["MASTER_PORT"] = "12355"  # Port number# This is important to prevent hangs or excessive memory utilization on GPU:0torch.cuda.set_device(rank)  # sets the default GPU for each processinit_process_group(backend="nccl", rank=rank, world_size=world_size)class Trainer:def __init__(self,model: torch.nn.Module,train_data: DataLoader,optimizer: torch.optim.Optimizer,gpu_id: int,save_every: int,) -> None:self.gpu_id = gpu_idself.train_data = train_dataself.optimizer = optimizerself.save_every = save_everyself.model = DDP(  # 感覺(jué)有點(diǎn)重復(fù), 上面 ddp_setup 已經(jīng)設(shè)置過(guò)默認(rèn) device 了model.to(gpu_id),  # 這里要先將模型放到 gpu_id 號(hào) GPU 上, 否則 DDP 會(huì)報(bào)錯(cuò)device_ids=[gpu_id],  # 那么這里再設(shè)置 device_ids 干嘛? 是可以分布到多個(gè) GPU 上嗎?)def _run_batch(self, source, targets):output = self.model(source)loss = F.cross_entropy(output, targets)self.optimizer.zero_grad()loss.backward()self.optimizer.step()def _run_epoch(self, epoch):b_sz = len(next(iter(self.train_data))[0])# len(self.train_data)} 將會(huì)被分割為 num_device 份print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")# sampler.set_epoch(epoch) is necessary to make shuffling work properly across multiple epochs.# Otherwise, the same ordering will be used in each epoch.self.train_data.sampler.set_epoch(epoch)  # 這里加了一句, 是為了保證每個(gè) epoch 的數(shù)據(jù)是隨機(jī)的for source, targets in self.train_data:source = source.to(self.gpu_id)targets = targets.to(self.gpu_id)self._run_batch(source, targets)def _save_checkpoint(self, epoch):ckp = self.model.module.state_dict()  # 因?yàn)?self.model 引用的是 DDP 對(duì)象, 所以想訪問(wèn)模型參數(shù), 則需要 .modulePATH = "checkpoint.pt"torch.save(ckp, PATH)print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")def train(self, max_epochs: int):for epoch in range(max_epochs):self._run_epoch(epoch)if self.gpu_id == 0 and epoch % self.save_every == 0:  # 主進(jìn)程才保存self._save_checkpoint(epoch)def load_train_objs():train_set = MyTrainDataset(2048)  # load your datasetmodel = torch.nn.Linear(20, 1)  # load your modeloptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)return train_set, model, optimizerdef prepare_dataloader(dataset: Dataset, batch_size: int):return DataLoader(dataset,batch_size=batch_size,pin_memory=True,shuffle=False,  # 有了 DistributedSampler, 這里就不用 shuffle 了, 不過(guò) default 已經(jīng)是 Falsesampler=DistributedSampler(dataset))def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):"""rank: Unique identifier of each process, GPU ID, 也是進(jìn)程的 ID, 0~world_size-1world_size: Total number of processes, 總共 GPU 數(shù)量"""ddp_setup(rank, world_size)  # 先設(shè)置當(dāng)前子進(jìn)程dataset, model, optimizer = load_train_objs()  # 之后似乎都一樣, 甚至數(shù)據(jù),模型,優(yōu)化器都是各進(jìn)程都創(chuàng)建train_data = prepare_dataloader(dataset, batch_size)trainer = Trainer(model, train_data, optimizer, rank, save_every)trainer.train(total_epochs)destroy_process_group()  # 結(jié)尾銷毀進(jìn)程組if __name__ == "__main__":import argparseparser = argparse.ArgumentParser(description='simple distributed training job')parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')parser.add_argument('save_every', type=int, help='How often to save a snapshot')parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')args = parser.parse_args()world_size = torch.cuda.device_count()# spawn processes, 自動(dòng)創(chuàng)建進(jìn)程, 并且把 rank 作為第一個(gè)參數(shù)傳入 mainmp.spawn(main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size)

更改代碼僅僅需要幾個(gè)步驟:

1. 構(gòu)建進(jìn)程組: init_process_group(...) & destroy_process_group()
main 函數(shù)被當(dāng)作子進(jìn)程啟動(dòng), 每個(gè)子進(jìn)程啟動(dòng)開(kāi)頭由 init_process_group(...) 構(gòu)建進(jìn)程組, 結(jié)尾由 destroy_process_group() 銷毀進(jìn)程組;
2. DistributedDataParallel 包裝模型
   其實(shí)主要還是持有參數(shù)的模型, 至于計(jì)算部分, 不要緊, 每個(gè)子進(jìn)程都在執(zhí)行相同的計(jì)算過(guò)程(除非設(shè)置了 if rank==... 的條件), 只會(huì)是參數(shù)梯度不同, 被包裝后的模型參數(shù)會(huì)自動(dòng)在進(jìn)程組之間同步;
   注意包裝前先將模型移動(dòng)到 GPU 上.
3. DistributedSampler 均勻地將樣本分給每個(gè)子進(jìn)程
   如果樣本數(shù)不夠整除, 則會(huì)將前幾個(gè)樣本補(bǔ)到末尾, 湊夠整除, 注意是打亂后的前幾個(gè), 相當(dāng)于隨機(jī)補(bǔ)幾個(gè)樣本;
   如果設(shè)置了 batch_size=32, 那么每個(gè)進(jìn)程都會(huì)得到 32 個(gè)樣本, 實(shí)際的 batch_size=32*num_gpus; 容易誤解的地方在于, 實(shí)際 batch_size 增大了, 那么我求 loss 時(shí)用 mean 的話, 會(huì)不會(huì)降低梯度大小? 不會(huì), 有些博主說(shuō)要 learning_rate*num_gpus, 但實(shí)際上人家的 ring all-reduce 算法是把各進(jìn)程上的梯度相加的, 相當(dāng)于執(zhí)行了多次梯度更新, 只不過(guò)是在相同的參數(shù)上, 而不是像單卡更新多次, 每次梯度計(jì)算在更新之后的不同的參數(shù)上.
   每個(gè)子進(jìn)程中訪問(wèn)的 DataLoader 中 batch 數(shù)會(huì)變?yōu)樵瓉?lái)的 1/num_gpus, 而 len(dataset) 不會(huì).
4. 每個(gè) epoch 開(kāi)始時(shí), 調(diào)用 train_loader.sampler.set_epoch(epoch), 否則, 將在每個(gè) epoch 中使用相同的順序.
5. 設(shè)置 if rank==0 為保存 checkpoint 的條件, 以保證只保存幾個(gè)相同模型的其中一個(gè).
   聚合操作, 如你想整合各進(jìn)程計(jì)算的不同結(jié)果并保存, 不應(yīng)在 if rank == 0 內(nèi), 聚合操作需要在每個(gè)進(jìn)程中執(zhí)行. 原因下面會(huì)解釋.
6. spawn 翻譯過(guò)來(lái)就是下蛋, 意思是啟動(dòng)子進(jìn)程, 可以看到, 相當(dāng)于執(zhí)行了多個(gè) main 函數(shù);
   rank 參數(shù)是自動(dòng)傳給 main 函數(shù)的.

BatchNorm 是根據(jù)數(shù)據(jù)計(jì)算均值和標(biāo)準(zhǔn)差的, 所以每個(gè) GPU 上計(jì)算的都不一樣, 如果想合成一個(gè)完整的大 Batch, 需要 SyncBatchNorm 同步.

1.4 multigpu_torchrun.py

import osimport torch
import torch.nn.functional as F
from torch import distributed
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSamplerfrom datautils import MyTrainDatasetdef ddp_setup():"""都不用設(shè)置主機(jī)地址和端口號(hào)了, 直接一個(gè) LOCAL_RANK"""torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))distributed.init_process_group(backend="nccl")class Trainer:def __init__(self,model: torch.nn.Module,train_data: DataLoader,optimizer: torch.optim.Optimizer,save_every: int,snapshot_path: str,) -> None:self.gpu_id = int(os.environ["LOCAL_RANK"])  # 這里也是自動(dòng)獲取 LOCAL_RANKself.model = model.to(self.gpu_id)self.train_data = train_dataself.optimizer = optimizerself.save_every = save_everyself.epochs_run = 0self.snapshot_path = snapshot_pathif os.path.exists(snapshot_path):print("Loading snapshot")self._load_snapshot(snapshot_path)self.model = DDP(self.model, device_ids=[self.gpu_id])def _load_snapshot(self, snapshot_path):loc = f"cuda:{self.gpu_id}"snapshot = torch.load(snapshot_path, map_location=loc, weights_only=True)  # 每個(gè) GPU 都要加載self.model.load_state_dict(snapshot["MODEL_STATE"])  # 之所以是 model.load_state_dict, 是因?yàn)樵?DDP 之前self.epochs_run = snapshot["EPOCHS_RUN"]print(f"Resuming training from snapshot at Epoch {self.epochs_run}")def _run_batch(self, source, targets):print(source.shape[0])output = self.model(source)loss = F.mse_loss(output, targets)self.optimizer.zero_grad()loss.backward()self.optimizer.step()return lossdef _run_epoch(self, epoch):b_sz = len(next(iter(self.train_data))[0])print(f"[GPU{self.gpu_id}] "f"Epoch {epoch} | "f"Batchsize: {b_sz} | "f"Steps: {len(self.train_data)} | "  # data_loader 會(huì) / num_devicesf"dsize: {len(self.train_data.dataset)}"  # 而數(shù)據(jù)集大小還是原來(lái)的)self.train_data.sampler.set_epoch(epoch)loss_epoch = 0for source, targets in self.train_data:source = source.to(self.gpu_id)targets = targets.to(self.gpu_id)loss = self._run_batch(source, targets)loss_epoch += lossprint(f"[GPU{self.gpu_id}] Loss {loss_epoch.item()}")distributed.all_reduce(loss_epoch, op=distributed.ReduceOp.AVG)print(f"[GPU{self.gpu_id}] Loss {loss_epoch.item()}")def _save_snapshot(self, epoch):snapshot = {"MODEL_STATE": self.model.module.state_dict(),  # 之后就要用 module 了"EPOCHS_RUN": epoch,}torch.save(snapshot, self.snapshot_path)print(f"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}")def train(self, max_epochs: int):for epoch in range(self.epochs_run, max_epochs):self._run_epoch(epoch)if self.gpu_id == 0 and epoch % self.save_every == 0:self._save_snapshot(epoch)def load_train_objs():train_set = MyTrainDataset(101)  # load your datasetmodel = torch.nn.Linear(20, 1)  # load your modeloptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)return train_set, model, optimizerdef prepare_dataloader(dataset: Dataset, batch_size: int):return DataLoader(dataset,batch_size=batch_size,pin_memory=True,shuffle=False,# 這個(gè) DistributedSampler 會(huì)自動(dòng)把數(shù)據(jù)集平均分給每個(gè) GPU, 只是每個(gè) DataLoader 得到的下標(biāo)是 len(dataset) / num_devices 個(gè)# 原來(lái)的 len(dataloder.dataset) 還是 len(dataset)# 注意會(huì)補(bǔ)全, 最后每個(gè) GPU 都會(huì)得到相同的數(shù)據(jù), 而不是最后一個(gè) GPU 會(huì)少得# 那補(bǔ)了之后, 樣本數(shù)是比源數(shù)據(jù)集多一些, 測(cè)試呢, 也就有偏差, 當(dāng)你有上萬(wàn)個(gè)測(cè)試樣本時(shí), 多出來(lái)的幾個(gè)樣本影響不大sampler=DistributedSampler(dataset))def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"):"""不帶 rank 了, 直接用 LOCAL_RANK"""ddp_setup()dataset, model, optimizer = load_train_objs()train_data = prepare_dataloader(dataset, batch_size)trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)trainer.train(total_epochs)distributed.destroy_process_group()if __name__ == "__main__":import argparseparser = argparse.ArgumentParser(description='simple distributed training job')parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')parser.add_argument('save_every', type=int, help='How often to save a snapshot')parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')args = parser.parse_args()main(args.save_every, args.total_epochs, args.batch_size)  # 不管 rank 和 device

執(zhí)行命令:

torchrun --standalone --nproc_per_node=2 multigpu_torchrun.py 50 10
# 如果設(shè)置 --nproc_per_node=gpu, 則自動(dòng)檢測(cè)可用 gpu 數(shù)量, 并為每個(gè) gpu 啟動(dòng)一個(gè)進(jìn)程.

這里使用了不同的啟動(dòng)方式 torchrun, 本質(zhì)還是一樣的, 特點(diǎn):

1. 能自動(dòng)重啟
   當(dāng)訓(xùn)練出現(xiàn)意外而中斷時(shí), torchrun 會(huì)自動(dòng)重啟, 如果保存了 checkpoint 并設(shè)置了自動(dòng)加載程序, 那么就可以接著訓(xùn).
2. 設(shè)置了環(huán)境變量 “LOCAL_RANK”
   你可以在程序中使用 os.environ["LOCAL_RANK"]) 訪問(wèn)當(dāng)前進(jìn)程的 rank 號(hào)了. 不過(guò)我感覺(jué)僅僅是在 distributed.init_process_group(backend="nccl") 之前使用, 后來(lái)的地方你可以接續(xù)這么干, 但構(gòu)建進(jìn)程組后有一個(gè)函數(shù) distributed.get_node_local_rank() 可以獲取進(jìn)程號(hào).
3. 單卡也可以跑, 設(shè)置 --nproc_per_node=1.

1.5 同步操作

模型參數(shù)可以通過(guò) DDP 自動(dòng)地同步, 那如果我想聚合所有子進(jìn)程上計(jì)算的 loss 呢? 或者我在測(cè)試時(shí), 想聚合測(cè)試結(jié)果? 官網(wǎng)的這個(gè)小教程沒(méi)教. 查閱博客才得知需要用 distributed.all_reduce(...).

上面的 multigpu_torchrun.py 中, 我已經(jīng)對(duì) loss 添加了這個(gè)同步:

print(f"[GPU{self.gpu_id}] Loss {loss_epoch.item()}")
distributed.all_reduce(loss_epoch, op=distributed.ReduceOp.AVG)
print(f"[GPU{self.gpu_id}] Loss {loss_epoch.item()}")
########## output ##########
[GPU1] Loss 0.3411558270454407
[GPU0] Loss 0.29943281412124634
[GPU1] Loss 0.3202943205833435
[GPU0] Loss 0.3202943205833435

兩個(gè) GPU 計(jì)算的 loss 分別為 0.3411558270454407 和 0.29943281412124634, 經(jīng)過(guò)同步, 都變?yōu)榱?0.3202943205833435.

注意:

• 只可對(duì) torch.Tensor 執(zhí)行同步, 其他類型的如 Python int 和 np.ndarray 都不行.
• 可以選擇其他聚合操作, 如 op=distributed.ReduceOp.SUM 表示相加:
  
  具體可見(jiàn): Collective Functions.
• 聚合操作不應(yīng)在 if rank == 0 內(nèi); 聚合操作需要在每個(gè)進(jìn)程中執(zhí)行.
  

1.6 會(huì)出現(xiàn)模型加載錯(cuò)誤

如果剛用 torch.save(...) 保存了模型, 立刻就使用 torch.load(...) 加載, 那么很可能會(huì)出現(xiàn)錯(cuò)誤:

[rank1]: RuntimeError: PytorchStreamReader failed reading zip archive: failed finding central directory

原因不明.

解決辦法:

time.sleep(1)
torch.load(...)

等 1s 再加載就不出錯(cuò)了.

2. 總結(jié)

看起來(lái)比較復(fù)雜, 但如果構(gòu)建對(duì) Distributed Data Parallel 的認(rèn)知框架, 一切都變得簡(jiǎn)單:

• DDP 為每個(gè) GPU 啟動(dòng)一個(gè)子進(jìn)程, 它們執(zhí)行"完全相同"的代碼;
• distributed.init_process_group(backend="nccl") 構(gòu)建進(jìn)程組, 程序結(jié)束時(shí) distributed.destroy_process_group() 銷毀進(jìn)程組;
• DistributedSampler(dataset) 為每個(gè)子進(jìn)程分發(fā)不重疊的等分的 Dataset 子集, 實(shí)現(xiàn)數(shù)據(jù)并行;
• 用 DDP 對(duì)象包裝模型, 就能在進(jìn)程組中同步梯度和參數(shù); 叫 ring all-reduce 算法;
• 你可以用 all_reduce 等操作實(shí)現(xiàn)進(jìn)程間的張量同步;
• 還可以根據(jù)進(jìn)程的 rank 號(hào)對(duì)不同子進(jìn)程執(zhí)行略有不同的操作, 如保存模型操作.