LLM大模型：deepspeed實戰和原理解析

第七子007發表於2024-07-28
大模型

　　多年前搞大資料，因為單節點無力儲存和計算PB級別的資料，所以hadoop這種分散式儲存和計算框架是標配！如今搞大模型，仍然需要對大量樣本資料做計算，因為涉及矩陣運算，單機單卡運算效率太低，也涉及到分散式計算了，大模型時代的分散式pre-train和Inference框架就有現成的—deepspeed！

　　1、老規矩，先直觀體驗一下deepspeed的使用：

　　 （1）自己定義一個簡單的模型：model.py

import torch
import numpy as np

class FashionModel(torch.nn.Module): 
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Linear(in_features=784, out_features=300),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=300, out_features=10)
        )

    def forward(self, batch_x):
        return self.seq(batch_x)

def img_transform(img): 
    img = np.asarray(img) / 255
    return torch.tensor(img, dtype=torch.float32).flatten()

　　（2）訓練的核心程式碼train.py:

import argparse
import torch
import torchvision
import deepspeed
from model import FashionModel, img_transform

# 命令列引數：deepspeed ds_train.py --epoch 5 --deepspeed --deepspeed_config ds_config.json
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=1, help="local rank passed from distributed launcher")
parser.add_argument("--epoch", type=int, default=1, help="epoch")
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()  # deepspeed命令列引數

dataset = torchvision.datasets.FashionMNIST(root='./dataset', download=True, transform=img_transform)  # 資料集
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True)  # 資料載入器，batch_size應該等於train_batch_size/gpu數量

model = FashionModel()  # 自定義的模型
model, _, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=model.parameters())  # deepspeed分散式模型
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(cmd_args.epoch):
    for x, y in dataloader:
        x = x.cuda()
        y = y.cuda()
        output = model(x)
        loss = loss_fn(output, y)
        model.backward(loss)  # 走deepspeed風格的backward
        model.step()
    print("epoch {} done...".format(epoch))
    model.save_checkpoint('./checkpoints')

　　配置檔案ds_config.json

{
    "train_batch_size": 128,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.00015
        }
    },
    "zero_optimization": {
        "stage": 2
    }
}

　　deepseed安裝好後，直接一行命令就開始執行：deepspeed ds_train.py --epoch 2 --deepspeed --deepspeed_config ds_config.json ；從日誌可以看出：有幾塊顯示卡就會生成幾個程序併發訓練；顯示卡之間使用nccl互相通訊；

　　主程序rank 0 列印日誌：

　　視訊記憶體都用上了：

　　訓練完畢生成的模型：

　　pre-train完成後的總要對LLM評估吧，程式碼ds_eval.py:

import argparse
from model import FashionModel, img_transform
import deepspeed
import torchvision
import torch

### deepspeed ds_eval.py --deepspeed --deepspeed_config ds_config.json
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()  # deepspeed命令列引數

model = FashionModel().cuda()  # 原始模型
model, optimizer, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=model.parameters())  # deepspeed分散式棋型
model.load_checkpoint('./checkpoints')  # 載入引數

model.eval()  # 分散式推理模式

# 只有主控程序帶頭做這些動作
if torch.distributed.get_rank() == 0:
    dataset = torchvision.datasets.FashionMNIST(root='./dataset', download=True, transform=img_transform)  # 訓練資料集
    batch_X = torch.stack([dataset[0][0], dataset[1][0]]).cuda()
    outputs = model(batch_X)  # 分散式推理
    print('分散式推理:', outputs.cpu().argmax(dim=1), [dataset[0][1], dataset[1][1]])

### 模型轉成torch單體
torch.save(model.module.state_dict(), 'model.pt')  # 儲存為普通torch模型引數
model = FashionModel().cuda()  # 載入torch模型
model.load_state_dict(torch.load('model.pt'))
model.eval()  # 單體推理
outputs = model(batch_X)
print('單體推理：', outputs.cpu().argmax(dim=1), [dataset[0][1], dataset[1][1]])

　　評估程式碼也簡單，也是直接一行命令：deepspeed ds_eval.py --deepspeed --deepspeed_config ds_config.json

import argparse
from model import FashionModel, img_transform
import deepspeed
import torchvision
import torch

### deepspeed ds_eval.py --deepspeed --deepspeed_config ds_config.json
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()  # deepspeed命令列引數

model = FashionModel().cuda()  # 原始模型
model, optimizer, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=model.parameters())  # deepspeed分散式棋型
model.load_checkpoint('./checkpoints')  # 載入引數

model.eval()  # 分散式推理模式

# 只有主控程序帶頭做這些動作
if torch.distributed.get_rank() == 0:
    dataset = torchvision.datasets.FashionMNIST(root='./dataset', download=True, transform=img_transform)  # 評估資料集
    batch_X = torch.stack([dataset[0][0], dataset[1][0]]).cuda()
    outputs = model(batch_X)  # 分散式推理
    print('分散式推理:', outputs.cpu().argmax(dim=1), [dataset[0][1], dataset[1][1]])

### 模型轉成torch單體
torch.save(model.module.state_dict(), 'model.pt')  # 儲存為普通torch模型引數
model = FashionModel().cuda()  # 載入torch模型
model.load_state_dict(torch.load('model.pt'))
model.eval()  # 單體推理
outputs = model(batch_X)
print('單體推理：', outputs.cpu().argmax(dim=1), [dataset[0][1], dataset[1][1]])

　　執行效果如下：

　　直觀感受完了deepspeed的使用，感覺比較簡單，底層分散式的訓練方案已經由框架都封裝好了，開發人員直接呼叫即可！ 接下來最重要的就是了解分散式訓練方案的原理了！

　　2、做分散式訓練，要麼是單節點視訊記憶體不夠，要麼是算力不夠。算力主要是各種矩陣運算啦，GPU硬體本身就是為這種計算定製的，軟體層面無法明顯最佳化，所以分散式系統主要最佳化的就是視訊記憶體的佔用啦！

　 （1）先看看單機單卡訓練時的視訊記憶體佔用，假設模型的引數量是m：

  • 引數本身儲存需要視訊記憶體，用FP32和FP16混合精度存放，需要6m視訊記憶體
  • 梯度儲存：用FP16存放，需要2m視訊記憶體
  • optimizer最佳化器：以adam為例，梯度下降的時候要存梯度和梯度平方，每個引數要存2個狀態，需要8m視訊記憶體

　　在不考慮存放訓練資料的前提下，pre-train至少需要6m+2m+8m=16m的視訊記憶體，所以後續的重點就是怎麼最佳化這三部分視訊記憶體佔用啦！

　（2）先來看看最簡單的一種情況：data parallel，簡稱DP。假設有N個顯示卡：

  • 就是把訓練資料均分成N份，然後N個顯示卡同時做forward和backward；N快顯示卡網路的初始引數都是一樣的
  • 產生的gradient都傳送給某個特定的顯示卡(這裡用0號顯示卡表示)
  • 0號顯示卡根據gradient更新自己網路的引數，然後把新的引數廣播傳送給其他所有顯示卡，讓所有顯示卡的網路引數保持一致
  • 除了0號顯示卡，其他顯示卡的作用就是計算loss和梯度

　　

　　這種DP方式的缺陷很明顯：0號顯示卡要收集其他所有顯示卡梯度，更新引數後要把新的引數廣播給所有顯示卡，顯示卡之間的通訊量很大，具體同步的資料量和顯示卡資料是線性正比的關係！

　（3）DP的改進版：distributed data parallel，簡稱DDP；和DP比，每塊卡單獨生成一個程序；

  • 資料照樣均分成，N個顯示卡同時做forward和backward；N快顯示卡網路的初始引數都是一樣的
  • 因為每張卡的資料不同，所以loss和梯度肯定不同，此時透過Ring-allReduce同步梯度，讓每張卡的梯度保持一致
  • 每張卡根據梯度更新自己的網路引數；因為每張卡的loss和梯度是要透過Ring-allReduce互相同步的，並且網路的初始狀態也是一樣的，所以每張卡的optomizer和網路狀態始終是一樣的！

　　顯示卡叢集總的資料傳送量：因為使用了Ring-allReduce傳輸資料（每個結點只給下一個結點傳輸資料，並不是整個叢集廣播），所以總的傳入傳出總量是固定的，不會因為顯示卡叢集擴大導致資料傳輸大增！

　　3、上述的DP和DDP，透過分散式增加了算力，但缺陷還是很明顯的：並未節約視訊記憶體！所以由此產生了ZeRO技術！

　 （1）預訓練時，optimizer佔用8倍引數量的視訊記憶體空間，是最耗費視訊記憶體的，所以肯定先從這種“大戶”下手啦！前面的DP和DDP，每塊顯示卡都儲存了完整的optimizer，互相都有冗餘，能不能消除這個冗餘了？比如叢集有3塊顯示卡，每塊顯示卡只存1/3的optimizer狀態？這就是ZeRO-1的思路！舉個例子：transformer不論decoer還是encoder，不是由一個個block上下疊加組成的麼？比如有12個block、3塊顯示卡，那麼每塊顯示卡儲存4個block的optimizer，不就ok啦？

  • 資料照樣均分成，N個顯示卡同時做forward和backward；N快顯示卡網路的初始引數都是一樣的
  • foward時所有顯示卡可以並行(因為都儲存和FP16的網路引數)，然後各自計算loss和梯度
  • 最關鍵的就是BP了：現在每塊顯示卡只存了部分optimizer，怎麼做BP更新引數了？
    • 因為每塊顯示卡都有完整的FP16網路引數，所以每塊顯示卡都可以並且需要根據loss計算梯度
    • 最後4個block的optimizer是GPU2負責，所以GPU0和1並不更新這4個block的引數。但是更新引數涉及梯度啊，GPU2的loss和梯度資訊不完整，這時就需要GPU0和1把自己計算的梯度資訊傳送給GPU2，整合後計算mean，用於更新最後4個block的引數！
    • 同理，中間4個block的梯度由GPU0和2傳送給CPU1，GPU1整合後計算mean，用於更新中間4個block的引數！最前面4個block的梯度由GPU1和1發給GPU0，GPU0整合後計算mean再更新網路引數！
    • 3塊顯示卡更新了各自負責block的網路引數，然後互相廣播，至此：每塊GPU的網路引數都是最新的了！

　　通訊量分析：和DDP是一樣的，但是每塊顯示卡節約了視訊記憶體！最核心的就是每塊顯示卡都把不屬於自己負責那段網路的梯度傳送給指定負責的網路卡，並未盲目全體廣播，此處節約了頻寬！但因為每塊顯示卡要廣播更新後的網路引數，所以網路通訊相比DDP並未減少！

　　這個思路有點像國內的鐵路局：國家在不同的區域分別設定了鐵路局，每個局負責自己片區鐵路線路的建設和運維；等建設完畢就把這個訊息傳送給其他鐵路局，然後開通相應的運輸路線！

　　（2）既然每塊GPU只負責更新部分引數，那是不是隻儲存對應的梯度也行了？這就是ZeRO-2的思路！

  • 資料照樣均分成，N個顯示卡同時做forward和backward；N快顯示卡網路的初始引數都是一樣的
  • foward時所有顯示卡可以並行(因為都儲存和FP16的網路引數)，然後各自計算loss和梯度
  • 做BP時：
    • GPU0和GPU1計算出最後4個block的梯度後發給GPU2，讓GPU2更新optimizer和網路引數，這部分的梯度自己都丟棄，完全不存；
    • 其他兩個block的引數做法類似，不再贅述
    • 最後3塊顯示卡再互相廣播更新後的網路引數

　（3）既然optimizer和梯度都可以只存部分，那引數是不是也可以了？這就是ZeRO-3的思路了！但這次的情況又有點不同：網路引數都不完整，怎麼forward？這就只能依靠廣播了，需要用到的時候讓其他GPU發過來！

  • 資料均分成N份，同時做forward；但因為每塊顯示卡的引數都不全，所以涉及到自己的時候要讓其他顯示卡發過來；比如最前面4個block做forward，GPU0有，但是GPU1和2沒有，就讓GPU0廣播；其他block同理，用的時候廣播，用完就丟棄不儲存！
  • BP計算loss和梯度也要網路引數啊，咋辦？同樣還是廣播的方式補全！

　　這種思路本質是需要用到的時候讓其他GPU配合傳送過來，用完就刪除不儲存！用顯示卡之間的頻寬換視訊記憶體的空間！通行量如下：

　　通行量是DDP的1.5倍，但是視訊記憶體佔用比DDP小了接近60倍！最後，來自ZeRO官方論文的總結對比：分別是DDP、ZeRO1/2/3階段的視訊記憶體消耗：

　　實戰中，一般採用ZeRO-2: 沒有增加通行量，但是極大減少了視訊記憶體的佔用！

其他

1、以前做大資料，hadoop是標配，會安裝、運維、調優甚至修改hadoop框架內部各種元件的研發很吃香，進大廠後工資都不低；同理，在以後AI時代，會安裝、運維、調優甚至更改分散式訓練/微調/推理框架的研發肯定更吃香，這絕對是個很不錯的方向！

2、顯示卡之間通訊，涉及到引數傳遞的，會讓顯示卡組成虛擬環，環內每個顯示卡的每個維度都依次給下一個顯示卡傳送資料，直到每個顯示卡的引數都一樣位置，這期間的經歷稱為scatter-reduce和all-gather！

　 scatter-reduce：單個維度向下擴散依次累加；這裡一看到reduce，就想起了10多年前因大資料爆火的map-reduce框架；這裡的reduce功能和map-reduce的功能原理上一模一樣，沒本質區別！

　　all-gather：單個完成的維度往下擴散，確保其他顯示卡該維度的資料是正確的！

　　在Ring-allReduce中，每塊顯示卡都在傳送和接受資料，可以最大程度利用每塊顯示卡的上下行頻寬！

參考：

1、https://www.bilibili.com/video/BV1LC4y1Y7tE/?spm_id_from=333.337.search-card.all.click&vd_source=241a5bcb1c13e6828e519dd1f78f35b2 deepspeed訓練和推理

2、https://www.bilibili.com/video/BV1fb421t7KN/?spm_id_from=333.337.search-card.all.click&vd_source=241a5bcb1c13e6828e519dd1f78f35b2 分散式大模型訓練

3、https://www.deepspeed.ai/ https://github.com/microsoft/DeepSpeed https://www.deepspeed.ai/getting-started/ 官網

4、https://www.bilibili.com/video/BV1ks4y1u7qr/?vd_source=241a5bcb1c13e6828e519dd1f78f35b2 DeepSpeed-Chat 模型訓練實戰

5、https://arxiv.org/pdf/1910.02054 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

6、https://www.bilibili.com/video/BV1mm42137X8/?spm_id_from=333.788.recommend_more_video.0&vd_source=241a5bcb1c13e6828e519dd1f78f35b2 DeepSpeed ZeRO技術

相關文章