DDP

Distributed Data Parallel 簡(jiǎn)稱 DDP，是 PyTorch 框架下一種適用于單機(jī)多卡、多機(jī)多卡任務(wù)的數(shù)據(jù)并行方式。由于其良好的執(zhí)行效率及廣泛的顯卡支持，熟練掌握 DDP 已經(jīng)成為深度學(xué)習(xí)從業(yè)者所必備的技能之一。本文結(jié)合具體代碼，詳細(xì)地說(shuō)明了 DDP 在項(xiàng)目中的使用方式。讀者按照本文所給的范例，只需稍經(jīng)調(diào)試，即可實(shí)現(xiàn) DDP 的整套流程。

概念辨析

具體講解 DDP 之前，我們先了解了解它和 Data Parallel (DP) 之間的區(qū)別。DP 同樣是 PyTorch 常見(jiàn)的多 GPU 并行方式之一，且它的實(shí)現(xiàn)非常簡(jiǎn)潔：

# 函數(shù)定義
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
'''
module : 模型
device_ids : 參與訓(xùn)練的 GPU 列表
output_device : 指定輸出的 GPU, 通常省略, 即默認(rèn)使用索引為 0 的顯卡
'''
# 程序模板
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

基本原理及固有缺陷：在 Data Parallel 模式下，數(shù)據(jù)會(huì)被自動(dòng)切分，加載到 GPU。同時(shí)，模型也將拷貝至各個(gè) GPU 進(jìn)行正向傳播。在多個(gè)進(jìn)程之間，會(huì)有一個(gè)進(jìn)程充當(dāng) master 節(jié)點(diǎn)，負(fù)責(zé)收集各張顯卡積累的梯度，并據(jù)此更新參數(shù)，再統(tǒng)一發(fā)送至其他顯卡。因此整體而言，master 節(jié)點(diǎn)承擔(dān)了更多的計(jì)算與通信任務(wù)，容易造成網(wǎng)絡(luò)堵塞，影響訓(xùn)練速度。

常見(jiàn)問(wèn)題及解決方案：Data Parallel 要求模型必須在 device_ids[0] 擁有參數(shù)及緩沖區(qū)，因此當(dāng)卡 0 被占用時(shí)，可以在 nn.DataParallel 之前添加如下代碼：

# 按照 PIC_BUS_ID 順序自 0 開(kāi)始排列 GPU 設(shè)備
os.environ['CUDA_DEVICE_ORDER'] = 'PIC_BUS_ID'
# 設(shè)置當(dāng)前使用的 GPU 為 2、3 號(hào)設(shè)備
os.environ['CUDA_VISIBLE_DEVICES'] = '2, 3'

如此，device_ids[0] 將被默認(rèn)為 2 號(hào)卡，device_ids[1] 則對(duì)應(yīng) 3 號(hào)卡

相較于 DP, Distributed Data Parallel 的實(shí)現(xiàn)要復(fù)雜得多，但是它的優(yōu)勢(shì)也非常明顯：

• DDP 速度更快，可以達(dá)到略低于顯卡數(shù)量的加速比；
• DDP 可以實(shí)現(xiàn)負(fù)載的均勻分配，克服了 DP 需要一個(gè)進(jìn)程充當(dāng) master 節(jié)點(diǎn)的固有缺陷；
• 采用 DDP 通常可以支持更大的 batch size，不會(huì)像 DP 那樣出現(xiàn)其他顯卡尚有余力，而卡 0 直接 out of memory 的情況；
• 另外，在 DDP 模式下，輸入到 data loader 的 bacth size 不再代表總數(shù)，而是每塊 GPU 各自負(fù)責(zé)的 sample 數(shù)量。比方說(shuō)，batch_size = 30，有兩塊 GPU。在 DP 模式下，每塊 GPU 會(huì)負(fù)責(zé) 15 個(gè)樣本。而在 DDP 模式下，每塊 GPU 會(huì)各自負(fù)責(zé) 30 個(gè)樣本；
• DDP 基本原理：倘若我們擁有 N 張顯卡，則在 Distributed Data Parallel 模式下，就會(huì)啟動(dòng) N 個(gè)進(jìn)程。每個(gè)進(jìn)程在各自的卡上加載模型，且模型的參數(shù)完全相同。訓(xùn)練過(guò)程中，各個(gè)進(jìn)程通過(guò)一種名為 Ring-Reduce 的方式與其他進(jìn)程通信，交換彼此的梯度，從而獲得所有的梯度信息。隨后，各個(gè)進(jìn)程利用梯度的平均值更新參數(shù)。由于初始值和更新量完全相同，所以各個(gè)進(jìn)程更新后的參數(shù)仍保持一致。

常用術(shù)語(yǔ)

• rank
  • 進(jìn)程號(hào)
  • 多進(jìn)程上下文中，通常假定 rank = 0 為主進(jìn)程或第一個(gè)進(jìn)程
• node
  • 物理節(jié)點(diǎn)，表示一個(gè)容器或一臺(tái)機(jī)器
  • 節(jié)點(diǎn)內(nèi)部可以包含多個(gè) GPU
• local_rank
  • 一個(gè) node 中，進(jìn)程的相對(duì)序號(hào)
  • local_rank 在 node 之間獨(dú)立
• world_size
  • 全局進(jìn)程數(shù)
  • 一個(gè)分布式任務(wù)中 rank 的數(shù)量
• group
  • 進(jìn)程組
  • 一個(gè)分布式任務(wù)就對(duì)應(yīng)一個(gè)進(jìn)程組
  • 只有當(dāng)用戶創(chuàng)立多個(gè)進(jìn)程組時(shí)，才會(huì)用到

代碼實(shí)現(xiàn)

Distributed Data Parallel 可以通過(guò) Python 的 torch.distributed.launch 啟動(dòng)器，在命令行分布式地執(zhí)行 Python 文件。執(zhí)行過(guò)程中，啟動(dòng)器會(huì)將當(dāng)前進(jìn)程（其實(shí)就是 GPU）的 index 通過(guò)參數(shù)傳遞給 Python，而我們可以利用如下方式獲取這個(gè) index：

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    metavar='N', help='Local process rank.')
args = parser.parse_args()
# print(args.local_rank)
# local_rank 表示本地進(jìn)程序號(hào)

隨后，初始化進(jìn)程組。對(duì)于在 GPU 執(zhí)行的任務(wù)，建議選擇 nccl (由 NVIDIA 推出) 作為通信后端。對(duì)于在 CPU 執(zhí)行的任務(wù)，建議選擇 gloo (由 Facebook 推出) 作為通信后端。倘若不傳入 init_method，則默認(rèn)為 env://，表示自環(huán)境變量讀取分布式信息

dist.init_process_group(backend='nccl', init_method='env://')
# 初始化進(jìn)程組之后, 通常會(huì)執(zhí)行這兩行代碼
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)
# 后續(xù)的 model = model.to(device), tensor.cuda(device)
# 對(duì)應(yīng)的都是這里由 args.local_rank 初始化得到的 device

數(shù)據(jù)部分，使用 Distributed Sampler 劃分?jǐn)?shù)據(jù)集，并將 sampler 傳入 data loader。需要注意的是，此時(shí)在 data loader 中不能指定 shuffle 為 True，否則會(huì)報(bào)錯(cuò) (sampler 已具備隨機(jī)打亂功能)

dev_sampler = data.DistributedSampler(dev_data_set)
train_sampler = data.DistributedSampler(train_data_set)
dev_loader = data.DataLoader(dev_data_set, batch_size=dev_batch_size, 
                             shuffle=False, sampler=dev_sampler)
train_loader = data.DataLoader(train_data_set, batch_size=train_batch_size, 
                               shuffle=False, sampler=train_sampler)

模型部分，首先將將模型送至 device，即對(duì)應(yīng)的 GPU 上，再使用 Distributed Data Parallel 包裝模型（順序顛倒會(huì)報(bào)錯(cuò)）

model = model.to(device)
model = nn.parallel.DistributedDataParallel(
    model, device_ids=[args.local_rank], output_device=args.local_rank
)

Distributed Data Parallel 模式下，保存模型應(yīng)使用 net.module.state_dict()，而非 net.state_dict()。且無(wú)論是保存模型，還是 LOGGER 打印，只對(duì) local_rank 為 0 的進(jìn)程操作即可，因此代碼中會(huì)有很多 args.local_rank == 0 的判斷

if args.local_rank == 0:
    LOGGER.info(f'saving latest model: {output_path}')
    torch.save({'model': model.module.state_dict(), 
                'optimizer': None, 'epoch': epoch, 'best-f1': best_f1}, 
               open(os.path.join(output_path, 'latest_model_{}.pth'.format(fold)), 'wb'))

利用 torch.load 加載模型時(shí)，設(shè)置 map_location=device，否則卡 0 會(huì)承擔(dān)更多的開(kāi)銷

load_model = torch.load(best_path, map_location=device)
model.load_state_dict(load_model['model'])

• dist.barrier() 可用于同步多個(gè)進(jìn)程，建議只在必要的位置使用，如初始化 DDP 模型之前、權(quán)重更新之后、開(kāi)啟新一輪 epoch 之前
• 計(jì)算 accuracy 時(shí)，可以使用 dist.all_reduce(score, op=dist.ReduceOp.SUM)，將各個(gè)進(jìn)程計(jì)算的準(zhǔn)確率求平均
• 計(jì)算 f1-score 時(shí)，可以使用 dist.all_gather(all_prediction_list, prediction_list)，將各個(gè)進(jìn)程獲得的預(yù)測(cè)值和真實(shí)值匯總到 all_list，再統(tǒng)一代入公式

啟動(dòng)方式

torch.distributed.launch

# 此處 --nproc_per_node 4 的含義是 server 有 4 張顯卡
python torch.distributed.launch --nproc_per_node 4 train.py
# 倘若使用 nohup, 則注意輸入命令后 exit 當(dāng)前終端
python torch.distributed.launch --nproc_per_node 4 train.py

• torchrun，推薦使用這種方式，因?yàn)?torch.distributed.launch 即將棄用

代碼中，只需將 Argument Parser 相關(guān)的部分替換為

local_rank = int(os.environ['LOCAL_RANK'])

然后將 args.local_rank 全部改為 local_rank 即可

啟動(dòng)命令

# 單機(jī)多卡訓(xùn)練時(shí), 可以不指定 nnodes
torchrun --nnodes=1 --nproc_per_node=4 train.py
# 倘若使用 nohup, 則注意輸入命令后 exit 當(dāng)前終端
nohup torchrun --nnodes=1 --nproc_per_node=4 train.py > nohup.out &

以上就是PyTorch Distributed Data Parallel使用詳解的詳細(xì)內(nèi)容，更多關(guān)于PyTorch Distributed Data Parallel的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

最新評(píng)論