單機多卡、多機多卡的藝術

2022 年 5 月 18 日
筆記
NLP通經脈

隨着接觸到的模型越來越大，自然就會接觸到這種技術。

記錄下自己的踩坑過程，當看到多機多卡跑通後，那種苦盡甘來的感覺還是挺舒服的。

我們首先來說一下單機多卡

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出現的進度條，會發現進度會同時是一樣的，然後就出現文章片頭出現的結果了。

Tags: NLP通經脈