隨著接觸到的模型越來越大，自然就會接觸到這種技術。
記錄下自己的踩坑過程，當看到多機多卡跑通後，那種苦盡甘來的感覺還是挺舒服的。
我們首先來說一下單機多卡
huggingface上面有大佬上傳了中文的BigBird的權重，想嘗試能夠處理的序列最長長度為4096的模型，但是放到單張卡里面batch_size基本上只能設定成2(16GB)，所以為了讓梯度下降更穩定，決定使用多卡進行訓練。本來是想嘗試把模型切成兩半，分別放到兩張卡里面，但是奈何自己沒有能力把bigbird轉換成nn.Sequential的樣子的型別，所以就放棄了，轉用DDP（Distributed Data Parallelism）。
（之後有關注了huggingface的Accelerate和另一個很有名氣的Colossal-AI,但是都會有同樣的bug）
我是參考這篇文章的：Distributed Training in PyTorch (Distributed Data Parallel) | by Praneet Bomma | Analytics Vidhya | Medium（良心文章，認真參考一次就跑通了）
現在來從頭開始，跑通單機多卡。
匯入依賴包
1 from distutils.command.config import config 2 import os 3 import jieba_fast 4 import json 5 import pandas as pd 6 import re 7 import torch 8 import numpy as np 9 import torch.nn.functional as F 10 import torch.optim as optim 11 import torch.nn as nn 12 import torch.distributed as dist 13 import torch.multiprocessing as mp 14 15 from tqdm.auto import tqdm 16 from transformers import BigBirdModel, BertTokenizer 17 from torch.utils.data import Dataset,DataLoader 18 from matplotlib import pyplot as plt 19 from datasets import load_dataset, load_metric 20 from torch.utils.tensorboard import SummaryWriter
編輯配置引數
1 class Config: 2 batch_size_train = 2 3 batch_size_valid = 1 4 5 max_length = 1500 6 seed = 4 7 device = torch.device("cuda:0") if torch.cuda.is_available() else 'cpu' 8 device1 = torch.device("cuda:1") if torch.cuda.is_available() else 'cpu' 9 # device = 'cpu' 10 bigbird_output_size = 768 11 vocab_size = 39999#+3 # len(tokenizer.get_vocab()) +3 是因為後面新增了特殊token 12 13 save_path = "model/BigBird_test3_v3_.bin" 14 15 epochs = 10 16 accumulate_setp = 10 17 18 gpus = 2 19 nr = 1 # global rank 第幾臺機器 20 nodes = 2 21 word_size = gpus*nodes
對於我來說，我不喜歡argument parser這種東西，所以我喜歡把配置引數放到一個類裡面：
對於單機多卡，真正要配置的只有最下面4個：
gpus: 一臺機器有多少張顯示卡
nr：number of rank 這裡指的是global rank，也就是在多機多卡環境下，每臺機器的編號，現在我們只有一臺機器，就設定為0。（多機多卡必須要有一個主機器，所以單機多卡是多機多卡，多機只有一臺機器的情況，主機器的global rank設定為0）
nodes:節點的個數（主機的臺數）
world_size:整個環境裡面，顯示卡的張數。
定義tokenizer和model
class JiebaTokenizer(BertTokenizer): ... class BB(torch.nn.Module): ...
自定義資料集
class DS(Dataset): ...
**定義train函式**
主要關注一下注釋部分，在自己的程式碼中新增需要新增的程式碼。
def train(gpu,config): rank = config.nr * config.gpus + gpu # train函式會執行到每個GPU上，所以需要顯示卡的ID 0~world_size-1 dist.init_process_group( backend='nccl', # 顯示卡的通訊方式 init_method='env://', # 初始化方法，從命令列的環境裡面讀取需要的環境變數 world_size=config.word_size, rank=rank ) torch.manual_seed(config.seed) # 設定隨機種子 tokenizer = JiebaTokenizer.from_pretrained('Lowin/chinese-bigbird-base-4096') model = BB() torch.cuda.set_device(gpu) # 選擇使用的GPU model.cuda(gpu) # 把模型放到被使用的GPU上 optimizer = optim.AdamW(params=model.parameters(),lr=1e-5,weight_decay=1e-2) model = nn.parallel.DistributedDataParallel(model,device_ids=[gpu],find_unused_parameters=True) # 需要把模型再次包裝成多GPU模型 trains = json.load(open("dataset/train.json")) dataSetTrain = DS(trains,tokenizer,config) train_sampler = torch.utils.data.distributed.DistributedSampler( dataSetTrain, num_replicas = config.word_size, rank = rank ) tDL = DataLoader(dataSetTrain,batch_size=config.batch_size_train,shuffle=False,pin_memory=True,sampler=train_sampler) step = 0 for epoch in range(config.epochs): if gpu == 0: # 第一張卡 (local rank) tDL = tqdm(tDL,leave=False) model.train() for batch in tDL: step += 1 labels = batch.pop('labels').cuda(non_blocking=True) # 把資料輸入輸出放到當前正在使用的顯示卡（編號為rank的那張顯示卡）裡面，non_blocking=True表示資料非同步載入到顯示卡里面 batch = {key:value.cuda(non_blocking=True) for key,value in batch.items()} logits = model(batch) loss_sum = F.cross_entropy(logits.view(-1,config.vocab_size),labels.view(-1),reduction='sum') # 下面三行是隻計算標題的梯度（任務是標題生成），進行梯度累計，可以不需要 title_length = labels.ne(0).sum().item() loss = loss_sum/title_length loss = loss/config.accumulate_setp loss.backward() if gpu == 0: # tqdm常用技巧，只讓GPU0上的模型的損失顯示出來（其他顯示卡的模型的損失是一樣的，為了不重複顯示，所以設定只讓0號GPU顯示結果） tDL.set_description(f'Epoch{epoch}') tDL.set_postfix(loss=loss.item()) if step % config.accumulate_setp == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 2) # 梯度裁剪，把梯度歸一化到01之間，讓梯度下降更穩定。 optimizer.step() optimizer.zero_grad()
#-----------------------------------------------------------------------下面的程式碼主要是儲存模型和驗證效能，可以不加--------------------------------------------------------------------------------------- if (epoch > 0) and (epoch % 2 == 0): torch.save(model.state_dict(), config.save_path+f'_epoch{epoch}') if ((gpu == 0) and (epoch % 2 == 0)) or epoch==(config.epochs-1): # 以下是評測驗證集的程式碼 tDL.write('*'*120) tDL.write(f'Epoch{epoch},開始評測效能') allIndexes = [] allLabels = [] with torch.no_grad(): model.eval() vDL = tqdm(vDL,leave=False) for sample in vDL: label = sample.pop('labels').cuda(non_blocking=True) sample = {key:value.cuda(non_blocking=True) for key,value in sample.items()} logits = model(sample) logits = logits[0] assert len(logits.shape) == 2 index = logits.argmax(dim=1) index = index>0 # 獲取token_id不為0的所有token 所在的輸出向量的索引 index = logits[index].argmax(dim=1) label = label[label!=0] allIndexes.append(index) allLabels.append(label) result = rouge.compute(predictions=allIndexes,references=allLabels) tDL.write(f'rouge1:{result["rouge1"][1][1]}') tDL.write(f'rouge2:{result["rouge2"][1][1]}') tDL.write(f'rougeL:{result["rougeL"][1][1]}')if gpu == 0: # 儲存最後一個epoch的模型 torch.save(model.state_dict(), config.save_path) writer.close()
定義main函式
def main(): config = Config() # 配置引數 os.environ['MASTER_ADDR'] = '10.100.132.151' # 主機器的IP，單機可以設定為localhost os.environ['MASTER_PORT'] = '12356' # 多機多卡時，不同機器和主機器之間的通訊埠，用於傳遞張量。 mp.spawn(train,nprocs=config.gpus,args=(config,)) # 開啟分散式訓練 train: 上面定義的train函式，nproc:每一臺機器有多少張顯示卡，args:配置引數 if __name__ == "__main__": main()
處理完成之後，就可以直接python xxx.py了，然後在終端輸入nvidia-smi後，會發現兩張卡都用起來了。
再來說一下多機多卡
搞定單機多卡後，多機多卡就只需要修改幾行程式碼，然後在不同的機器上分別啟動就好了。
只需要修改配置引數，就可以實現多級多卡了：
class Config: ... nr = 0 # global rank 第幾臺機器，0表示主機器 nodes = 2 # 把這裡修改為2，表示我有2臺機器 word_size = gpus*nodes
然後再到另外一臺機器上，也修改config引數：
class Config: ... nr = 1 # global rank 第二臺機器，0表示第一臺機器 nodes = 2 # 也把這裡修改為2 word_size = gpus*nodes
然後分別在兩臺主機上使用python xxx.py，對比兩臺機器的tqdm出現的進度條，會發現進度會同時是一樣的，然後就出現文章片頭出現的結果了。