本文分享自華為雲社群《大模型LLM之分散式訓練》，作者： 碼上開花_Lancer。
隨著語言模型引數量和所需訓練資料量的急速增長，單個機器上有限的資源已無法滿足大語言模型訓練的要求。需要設計分散式訓練（Distributed Training）系統來解決海量的計算和記憶體資源要求問題。
在分散式訓練系統環境下需要將一個模型訓練任務拆分成多個子任務，並將子任務分發給多個計算裝置，從而解決資源瓶頸。但是如何才能利用包括數萬計算加速晶片的叢集，訓練模型引數量千億甚至是萬億的大規模語言模型？這其中涉及到叢集架構、並行策略、模型架構、記憶體最佳化、計算最佳化等一系列的技術。
我將詳細介紹分散式機器學習系統的基礎概念、分散式訓練叢集架構、分散式訓練並行策略，並以DeepSpeed 為例介紹如何在叢集上訓練大語言模型。
一、分散式訓練概述
分散式訓練（Distributed Training）是指將機器學習或深度學習模型訓練任務分解成多個子任務，並在多個計算裝置上並行地進行訓練。圖1給出了單個計算裝置和多個計算裝置的示例，這裡計算裝置可以是中央處理器（Central Processing Unit，CPU）、圖形處理器（Graphics Processing Unit，GPU）、張量處理器（Tensor Processing Unit，TPU）也可以是神經網路處理器（Neural network Processing Unit，NPU）。
由於同一個伺服器內部的多個計算裝置之間記憶體也可能並不共享，因此無論這些計算裝置是否處於一個伺服器還是多個伺服器中，其系統架構都屬於分散式系統範疇。一個模型訓練任務往往會有大量的訓練樣本作為輸入，可以利用一個計算裝置完成，也可以將整個模型的訓練任務拆分成子任務，分發給不同的計算裝置，實現平行計算。
此後，還需要對每個計算裝置的輸出進行合併，最終得到與單個計算裝置等價的計算結果。由於每個計算裝置只需要負責子任務，並且多個計算裝置可以並行執行，因此其可以更快速地完成整體計算，並最終實現對整個計算過程的加速。
圖1 單計算裝置計算和多計算裝置示例
促使人們設計分散式訓練系統的一個最重要的原因就是單個計算裝置的算力已經不足以支撐模型訓練。圖2給出了機器學習模型對於算力的需求以及同期單個計算裝置能夠提供的算力。如圖所示，機器學習模型快速發展，從2013 年AlexNet 開始，到2022 年擁有5400 億引數的PalM 模型，機器學習模型以每18 個月增長56 倍的速度發展。模型引數規模增大的同時，對訓練資料量的要求也指數級增長，這更加劇了對算力的需求。
然而，近幾年CPU 的算力增加已經遠低於摩爾定律（Moore’s Law），雖然計算加速裝置（如GPU、TPU 等）為機器學習模型提供了大量的算力，但是其增長速度仍然沒有突破每18 個月翻倍的摩爾定律。為了能夠滿足機器學習模型的發展，只有透過分散式訓練系統才可以匹配模型不斷增長的算力需求。
圖2 機器學習模型引數量增長和計算硬體的算力增長對比
分散式訓練的總體目標就是提升總的訓練速度，減少模型訓練的總體時間。總訓練速度可以用如下公式簡略估計：
總訓練速度∝ 單裝置計算速度× 計算裝置總量× 多裝置加速比
其中，單裝置計算速度主要由單塊計算加速晶片的運算速度和資料I/O 能力來決定，對單裝置訓練效率進行最佳化，主要的技術手段有混合精度訓練、運算元融合、梯度累加等；分散式訓練系統中計算裝置數量越多，其理論峰值計算速度就會越高，但是受到通訊效率的影響，計算裝置數量增大則會造成加速比急速降低；多裝置加速比則是由計算和通訊效率決定，需要結合演算法和網路拓撲結構進行最佳化，分散式訓練並行策略主要目標就是提升分散式訓練系統中的多裝置加速比。
大語言模型引數量和所使用的資料量都非常巨大，因此都採用了分散式訓練架構完成訓練。文獻[5] 針對GPT-3 的訓練過程僅介紹了訓練過程全部使用NVIDIA V100 GPU，文獻[31] 介紹了OPT 使用了992 塊NVIDIA A100 80G GPU，採用全分片資料並行（Fully Shared Data Parallel）[129]以及Megatron-LM 張量並行（Tensor Parallelism）[130]，整體訓練時間將近2 個月。
BLOOM[33] 模型的研究人員則公開了更多在硬體和所採用的系統架構方面的細節。該模型的訓練一共花費3.5 個月，使用48 個計算節點。每個節點包含8 塊NVIDIA A100 80G GPU（總計384 個GPU），並且使用4*NVLink 用於節點內部GPU 之間通訊。節點之間採用四個Omni-Path 100 Gbps 網路卡構建的增強8 維超立方體全域性拓撲網路進行通訊。
文獻[37] 並沒有給出LLaMA 模型訓練中所使用的叢集的具體配置和網路拓撲結構，但是給出了不同引數規模的總GPU 小時數。LLaMA 模型訓練採用A100-80GB GPU，LLaMA-7B 模型訓練需要82432 GPU 小時，LLaMA-13B 模型訓練需要135168GPU 小時，LLaMA-33B 模型訓練花費了530432 GPU 小時，而LLaMA-65B 模型訓練花費則高達1022362 GPU 小時。由於LLaMA 所使用的訓練資料量遠超OPT 和BLOOM 模型，因此，雖然模型引數量遠小於上述兩個模型，但是其所需計算量仍然非常驚人。
透過使用分散式訓練系統，大語言模型訓練週期可以從單計算裝置花費幾十年，縮短到使用數千個計算裝置花費幾十天就可以完成。然而，分散式訓練系統仍然需要克服計算牆、視訊記憶體牆、通訊牆等多種挑戰，以確保叢集內的所有資源得到充分利用，從而加速訓練過程並縮短訓練週期。
• 計算牆：單個計算裝置所能提供的計算能力與大語言模型所需的總計算量之間存在巨大差異。2022 年3 年釋出的NVIDIA H100 SXM 的單卡FP16 算力也只有2000 TFLOPs，而GPT-3
則需要314 ZFLOPs 的總算力，兩者相差了8 個數量級。
• 視訊記憶體牆：單個計算裝置無法完整儲存一個大語言模型的引數。GPT-3 包含1750 億引數，如果採用FP16 格式進行儲存，需要700GB 的計算裝置記憶體空間，而NVIDIA H100 GPU 只有80 GB 視訊記憶體。
• 通訊牆：分散式訓練系統中各計算裝置之間需要頻繁地進行引數傳輸和同步。由於通訊的延遲和頻寬限制，這可能成為訓練過程的瓶頸。GPT-3 訓練過程中，如果分散式系統中存在128個模型副本，那麼在每次迭代過程中至少需要傳輸89.6TB 的梯度資料。而截止2023 年8 月，單個InfiniBand 鏈路僅能夠提供不超過800Gb/s 頻寬。計算牆和視訊記憶體牆源於單計算裝置的計算和儲存能力有限，與模型對龐大計算和儲存需求之間存在矛盾。這個問題可以透過採用分散式訓練方法來解決，但分散式訓練又會面臨通訊牆的挑戰。在多機多卡的訓練中，這些問題逐漸顯現。隨著大模型引數的增大，對應的叢集規模也隨之增加，這些問題變得更加突出。同時，在大型叢集進行長時間訓練時，裝置故障可能會影響或中斷訓練過程，對分散式系統的問題性也提出了很高要求。
二、分散式訓練並行策略
分散式訓練系統目標就是將單節點模型訓練轉換成等價的分散式並行模型訓練。對於大語言模型來說，訓練過程就是根據資料和損失函式，利用最佳化演算法對神經網路模型引數進行更新的過程。單節點模型訓練系統結構如圖3所示，主要由資料和模型兩個部分組成。訓練過程會由多個資料小批次（Mini-batch）完成。
圖中資料表示一個資料小批次。訓練系統會利用資料小批次根據損失函式和最佳化演算法生成梯度，從而對模型引數進行修正。針對大語言模型多層神經網路的執行過程，可以由一個計算圖（Computational Graph）表示。這個圖有多個相互連線的運算元（Operator），每個運算元實現一個神經網路層（Neural Network Layer），而引數則代表了這個層在訓練中所更新的權重。
圖3 單裝置模型訓練系統
計算圖的執行過程可以分為前向計算和反向計算兩個階段。前向計算的過程是將資料讀入第一個運算元，計算出相應的輸出結構，然後依此重複這個前向計算過程，直到最後一個運算元結束。反向計算過程，是根據最佳化函式和損失，每個運算元依次計算出梯度，並利用梯度更新本地的引數。在反向計算結束後，該資料小批次的計算完成，系統就會讀取下一個資料小批次，繼續下一輪的模型引數更新。
根據單裝置模型訓練系統的流程，可以看到如果進行並行加速，可以從資料和模型兩個維度進行考慮。首先可以對資料進行切分（Partition），並將同一個模型複製到多個裝置上，並行執行不同的資料分片，這種方式通常被稱為資料並行（Data Parallelism，DP）。還可以對模型進行劃分，將模型中的運算元分發到多個裝置分別完成，這種方式通常被稱為模型並行（Model Parallelism，MP）。當訓練超大規模語言模型時，往往需要同時對資料和模型進行切分，從而實現更高程度的並行，這種方式通常被稱為混合並行（Hybrid Parallelism，HP）。
2.1、資料並行
在資料並行系統中，每個計算裝置都有整個神經網路模型的完整副本（Model Replica），進行迭代時，每個計算裝置只分配了一個批次資料樣本的子集，並根據該批次樣本子集的資料進行網路模型的前向計算。假設一個批次的訓練樣本數為N，使用M 個計算裝置平行計算，每個計算裝置會分配到N/M 個樣本。前向計算完成後，每個計算裝置都會根據本地樣本計算損失誤差得到梯度Gi（i 為加速卡編號），並將本地梯度Gi 進行廣播。所有計算裝置需要聚合其他加速度卡給出的梯度值，然後使用平均梯度(ΣNi=1Gi)/N 對模型進行更新，完成該批次訓練。圖4給出了由兩個計算裝置組成的資料並行訓練系統樣例。
圖4 兩節點資料並行訓練系統樣例
資料並行訓練系統可以透過增加計算裝置，有效提升整體訓練吞吐量，每秒全域性批次數（Global Batch Size Per Second) 。它和單計算裝置訓練相比，最主要的區別就在於反向計算中的梯度需要在所有計算裝置中進行同步，以保證每個計算裝置上最終得到的是所有程序上梯度的平均值。
常見的神經網路框架中都有資料並行方式的具體實現，包括：TensorFlow DistributedStrategy、PyTorch Distributed、Horovod DistributedOptimizer 等。由於基於Transformer 架構的大語言模型中每個運算元都是依賴單個資料而非批次資料，因此資料並行並不會影響其計算邏輯，一般情況下各訓練裝置中前向計算是獨立的，不涉及同步問題。資料並行訓練加速比最高，但要求每個裝置上都備份一份模型，視訊記憶體佔用比較高。
使用PyTorch DistributedDataParallel 實現單個伺服器多加速卡訓練程式碼如下，首先構造DistributedSampler類，將資料集的樣本隨機打亂並分配到不同計算裝置：
class DistributedSampler(Sampler): def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True, seed=0): if num_replicas is None: if not dist.is_available(): raise RuntimeError("Requires distributed package to be available") num_replicas = dist.get_world_size() if rank is None: if not dist.is_available(): raise RuntimeError("Requires distributed package to be available") rank = dist.get_rank() self.dataset = dataset # 資料集 self.num_replicas = num_replicas # 程序個數預設等於world_size(GPU 個數) self.rank = rank # 當前屬於哪個程序/哪塊GPU self.epoch = 0 self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) # 每個程序的樣本個數 self.total_size = self.num_samples * self.num_replicas # 資料集總樣本的個數 self.shuffle = shuffle # 是否要打亂資料集 self.seed = seed def __iter__(self): # 1、Shuffle 處理：打亂資料集順序 if self.shuffle: # 根據epoch 和種子進行混淆 g = torch.Generator() # 這裡self.seed 是一個定值，透過set_epoch 改變self.epoch 可以改變我們的初始化種子 # 這就可以讓每一個epoch 中資料集的打亂順序不同，使每一個epoch 中， # 每一塊GPU 拿到的資料都不一樣，這樣可以有利於更好的訓練 g.manual_seed(self.seed + self.epoch) indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) # 資料補充 indices += indices[:(self.total_size - len(indices))] assert len(indices) == self.total_size # 分配資料 indices = indices[self.rank:self.total_size:self.num_replicas] assert len(indices) == self.num_samples return iter(indices) def __len__(self): return self.num_samples def set_epoch(self, epoch): self.epoch = epoch
利用DistributedSampler 構造完整的訓練程式樣例main.py 如下：
import argparse import os import shutil import time import warnings import numpy as np warnings.filterwarnings('ignore') import torch import torch.nn as nn import torch.nn.parallel import torch.backends.cudnn as cudnn import torch.distributed as dist import torch.optim import torch.utils.data import torch.utils.data.distributed from torch.utils.data.distributed import DistributedSampler from models import DeepLab from dataset import Cityscaples parser = argparse.ArgumentParser(description='DeepLab') parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument('--epochs', default=100, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)') parser.add_argument('-b', '--batch-size', default=3, type=int, metavar='N') parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training') args = parser.parse_args() torch.distributed.init_process_group(backend="nccl") # 初始化 print("Use GPU: {} for training".format(args.local_rank)) # create model model = DeepLab() torch.cuda.set_device(args.local_rank) # 當前顯示卡 model = model.cuda() # 模型放在顯示卡上 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) # 資料並行 criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) train_dataset = Cityscaples() train_sampler = DistributedSampler(train_dataset) # 分配資料 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler)
透過以下命令列啟動上述程式：
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main.py
2.2 模型並行
模型並行（Model Parallelism）往往用於解決單節點記憶體不足的問題。以包含1750 億引數的GPT-3 模型為例，如果模型中每一個引數都使用32 位浮點數表示，那麼模型需要佔用700GB（即175G× 4 Bytes）記憶體，如果使用16 位浮點表示，每個模型副本需要也需要佔用350GB 記憶體。以2022 年3 月NVIDIA 釋出的H100 加速卡也僅支援80GB 視訊記憶體，無法將整個模型完整放入其中。模型並行可以從計算圖角度，以下兩種形式進行切分：
（1）按模型的層切分到不同裝置，即層間並行或運算元間並行（Inter-operator Parallelism），也稱之為流水線並行（Pipeline Parallelism，PP）；
（2）將計算圖層內的引數切分到不同裝置，即層內並行或運算元內並行（Intra-operator Parallelism），也稱之為張量並行（Tensor Parallelism，TP）。
兩節點模型並行訓練系統樣例如圖4.9所示，左邊為流水線並行，模型的不同層被切分到不同的裝置中；右邊為張量並行，同一個層中的不同的引數被切分到不同的裝置中進行計算。
流水線並行
流水線並行（Pipeline Parallelism，PP）是一種平行計算策略，將模型的各個層分段處理，並將每個段分佈在不同的計算裝置上，使得前後階段能夠流水式、分批進行工作。流水線並行通常應用於大規模模型的並行系統中，以有效解決單個計算裝置記憶體不足的問題。圖4.6給出了一個由四個計算裝置組成的流水線並行系統，包含了前向計算和後向計算。其中F1、F2、F3、F4 分別代表四個前向路徑，位於不同的裝置上；而B4、B3、B2、B1 則代表逆序的後向路徑，也分別位於四個不同的裝置上。然而，從圖中可以看出，計算圖中的下游裝置（Downstream Device）需要長時間持續處於空閒狀態，等待上游裝置（Upstream Device）的計算完成，才能開始計算自身的任務。
圖5 兩節點模型並行訓練系統樣例
這種情況導致了裝置的平均使用率大幅降低，形成了模型並行氣泡（Model Parallelism Bubble），也稱為流水線氣泡（Pipeline Bubble）。
圖6 流水線並行樣例
樸素流水線策略所產生的並行氣泡，使得系統無法充分利用計算資源，降低了系統整體的計算效率。為了能夠減少並行氣泡，文獻[131] 提出了GPipe 方法，將小批次（Mini-batch）進一步劃分成更小的微批次（Micro-batch），利用流水線並行方案，每次處理一個微批次的資料。
在當前階段計算完成得到結果後，將該微批次的結果傳送給下游裝置，同時開始處理後一個微批次的資料，這樣可以在一定程度上減少並行氣泡。圖7GPipe 策略流水線並行樣例。如圖所示，前向F1計算被拆解為了F11，F12，F13，F14，在計算裝置1 中計算完成F11 後，會在計算裝置2 中開始進行F21 計算，同時計算裝置1 中並行開始F12 的計算。相比於最原始的流水線並行方法，GPipe 流水線方法可以有效降低並行氣泡。
圖7 GPipe 策略流水線並行樣例
GPipe 策略雖然可以減少一定的並行氣泡，但是隻有當一個Mini-batch 中所有的前向計算完成後，才能開始執行後向計算。因此還是會產生很多並行氣泡，從而降低了系統的並行效率。Megatron-LM[132] 提出了1F1B 流水線策略，即一個前向通道和一個後向通道。1F1B 流水線策略引入了任務排程機制，使得下游裝置能夠在等待上游計算的同時執行其他可並行的任務，從而提高裝置的利用率。1F1B 給出了非交錯式和交錯式兩種方式排程方式，如圖8所示。
1F1B 非交錯式排程模式可分為三個階段。首先是熱身階段，在該階段中，計算裝置中進行不同數量的前向計算。接下來的階段是前向-後向階段，計算裝置按順序執行一次前向計算，然後進行一次後向計算。最後一個階段是後向階段，計算裝置在完成最後一次後向計算。相比於GPipe 策略，非交錯式排程模式在節省記憶體方面表現更好。然而，它需要與GPipe 策略一樣的時間來完成一輪計算。
1F1B 交錯式排程模式要求micro-batch 的數量是流水線階段的整數倍。每個裝置不再僅負責連續多個層的計算，而是可以處理多個層的子集，這些子集被稱為模型塊。具體而言，在之前的模式中，裝置1 可能負責層1-4，裝置2 負責層5-8，以此類推。然而，在新的模式下，裝置1 可以處理層1、2、9、10，裝置2 處理層3、4、11、12，以此類推。這種模式下，每個裝置在流水線中被分配到多個階段。例如，裝置1 可能參與熱身階段、前向計算階段和後向計算階段的某些子集任務。每個裝置可以並行執行不同階段的計算任務，從而更好地利用流水線並行的優勢。這種模式不僅在記憶體消耗方面表現出色，還能夠提高計算效率，使得大型模型的並行系統能夠更高效地完成計算任務。
圖8 1F1B 流水線並行策略樣例
PyTorch 中也包含了實現流水線的API 函式Pipe，具體實現參考“torch.distributed.pipeline.sync.Pipe”類。可以使用這個API 構造一個包含兩個線性層，分別放置在2 個不同計算裝置中的樣例如下：
{# Step 0. Need to initialize RPC framework first. os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '29500' torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1) # Step 1: build a model including two linear layers fc1 = nn.Linear(16, 8).cuda(0) fc2 = nn.Linear(8, 4).cuda(1) # Step 2: wrap the two layers with nn.Sequential model = nn.Sequential(fc1, fc2) # Step 3: build Pipe (torch.distributed.pipeline.sync.Pipe) model = Pipe(model, chunks=8) # do training/inference input = torch.rand(16, 16).cuda(0) output_rref = model(input) }
張量並行
張量並行（Tensor Parallelism，TP）需要根據模型的具體結構和運算元型別，解決如何將引數切分到不同裝置，以及如何保證切分後數學一致性兩個問題。大語言模型都是以Transformer 結構為基礎，Transformer 結構主要由以下三種運算元構成：嵌入式表示（Embedding）、矩陣乘（MatMul）和交叉熵損失（Cross Entropy Loss）計算構成。
這三種型別的運算元有較大的差異，都需要設計對應的張量並行策略[130]，才可以實現將引數切分到不同的裝置。對於嵌入表示（Embedding）運算元，如果總的詞表數非常大，會導致單計算裝置視訊記憶體無法容納Embedding 層引數。舉例來說，如果詞表數量是64000，嵌入表示維度為5120，型別採用32 位精度浮點數，那麼整層引數需要的視訊記憶體大約為64000 × 5120 × 4/1024/1024 = 1250MB，反向梯度同樣需要1250MB，僅僅儲存就需要將近2.5GB。
對於嵌入表示層的引數，可以按照詞維度切分，每個計算裝置只儲存部分詞向量，然後透過彙總各個裝置上的部分詞向量，從而得到完整的詞向量。圖4.9給出了單節點Embedding 和兩節點張量並行的示意圖。
在單節點上，執行Embedding 操作，bz 是批次大小（batch size），Embedding 的引數大小為[word_size, hidden_size]，計算得到[bz,hidden_size] 張量。圖4.9中Embedding 張量並行示例將Embedding 引數沿word_size 維度，切分為兩塊，每塊大小為[word_size/2, hidden_size]，分別儲存在兩個裝置上。當每個節點查詢各自的詞表時，如果無法查到，則該詞的表示為0，各自裝置查詢後得到[bz, hidden_size] 結果張量，最後透過AllReduce_Sum 通訊¬，跨裝置求和，得到完整的全量結果，可以看出，這裡的輸出結果和單計算裝置執行的結果一致。
圖9 兩節點Embedding 運算元張量並行示例
矩陣乘（MatMul）的張量並行要充分利用矩陣了分塊乘法原理。舉例來說，要實現如下矩陣乘法Y = X ×A，其中X 是維度為M × N 的輸入矩陣，A 是維度為N ×K 的引數矩陣，Y 是結果矩陣，維度為M ×K。如果引數矩陣A 非常大，甚至超出單張卡的視訊記憶體容量，那麼可以把引數矩陣A 切分到多張卡上，並透過集合通訊匯集結果，保證最終結果在數學計算上等價於單計算裝置計算結果。引數矩陣A 存在兩種切分方式：
(1) 引數矩陣A 按列切塊，將矩陣A 按列切成：A = [A1,A2]
(2) 引數矩陣A 按行切塊，將矩陣A 按行切成：
圖10給出了引數矩陣按列切分的示例，引數矩陣A 分別將A1，A2 放置在兩個計算裝置上。兩個計算裝置分別計算Y1 = X ×A1 和Y2 = X ×A2。計算完成後，多計算裝置間進行通訊，從而獲取其它計算裝置上的計算結果，並拼接在一起得到最終的結果矩陣Y ，該結果在數學上與單計算裝置計算結果上完全等價。
圖10 兩節點矩陣乘運算元張量並行按列切分示例
圖11給出了引數矩陣按列行分的示例，為了滿足矩陣乘法規則，輸入矩陣X 需要按列切分X = [X1|X2]。同時，將矩陣分塊，分別放置在兩個計算裝置上，每個計算裝置分別計算Y1 =X1 ×A1 和Y2 = X2 ×A2。計算完成後，多個計算裝置間通訊獲取歸約其他卡上的計算結果，可以得到最終的結果矩陣Y 。同樣，這種切分方式，既可以保證數學上的計算等價性，並解決單計算裝置視訊記憶體無法容納，又可以保證單計算裝置透過拆分方式可以裝下引數A 的問題。
Transformer 中的FFN 結構均包含兩層全連線（FC）層，即存在兩個矩陣乘，這兩個矩陣乘分別採用上述兩種切分方式，如圖4.12所示。對第一個FC 層的引數矩陣按列切塊，對第二個FC層引數矩陣按行切塊。這樣第一個FC 層的輸出恰好滿足第二個FC 層資料輸入要求（按列切分），因此可以省去第一個FC 層後的彙總通訊操作。多頭自注意力機制的張量並行與FFN 類似，因為具有多個獨立的頭，因此相較於FFN 更容易實現並行，其矩陣切分方式如圖4.13所示。具體可以參考文獻[130]。
分類網路最後一層一般會選用Softmax 和Cross_entropy 運算元來計算交叉熵損失（Cross Entropy Loss）。如果類別數量非常大，會導致單計算裝置記憶體無法儲存和計算logit 矩陣。針對這一類運算元，可以按照類別維度切分，同時透過中間結果通訊，得到最終的全域性的交叉熵損失。
圖11 兩節點矩陣乘運算元張量並行按行切分示例
圖12 FNN 結構張量並行示意圖
首先計算的是softmax 值，公式如下：
其中，p 表示張量並行的裝置號。得到Softmax 計算結果之後，同時對標籤Target 按類別切分，每個裝置得到部分損失，最後再進行一次通訊，得到所有類別的損失。整個過程，只需要進行三次小量的通訊，就可以完成交叉熵損失的計算。PyTorch 提供了細粒度張量級別的並行API，DistributedTensor。也提供了粗粒度模型層面的API 對“nn.Module”進行張量並行。透過以下幾行程式碼就可以實現對一個大的張量進行分片：
import torch from torch.distributed._tensor import DTensor, DeviceMesh, Shard, distribute_tensor # construct a device mesh with available devices (multi-host or single host) device_mesh = DeviceMesh("cuda", [0, 1, 2, 3]) # if we want to do row-wise sharding rowwise_placement=[Shard(0)] # if we want to do col-wise sharding colwise_placement=[Shard(1)] big_tensor = torch.randn(888, 12) # distributed tensor returned will be sharded across the dimension specified in placements rowwise_tensor = distribute_tensor(big_tensor, device_mesh=device_mesh, placements=rowwise_placement)
對於像“nn.Linear”這樣已經有“torch.Tensor”作為引數的模組，也提供了模組級API “distribute_module”在模型層面進行張量並行，參考程式碼如下：
import torch from torch.distributed._tensor import DeviceMesh, Shard, distribute_tensor,distribute_module class MyModule(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(8, 8) self.fc2 = nn.Linear(8, 8) self.relu = nn.ReLU() def forward(self, input): return self.relu(self.fc1(input) + self.fc2(input)) mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1], [2, 3]]) def shard_params(mod_name, mod, mesh): rowwise_placement = [Shard(0)] def to_dist_tensor(t): return distribute_tensor(t, mesh, rowwise_placement) mod._apply(to_dist_tensor) sharded_module = distribute_module(MyModule(), mesh, partition_fn=shard_params) def shard_fc(mod_name, mod, mesh): rowwise_placement = [Shard(0)] if mod_name == "fc1": mod.weight = torch.nn.Parameter(distribute_tensor(mod.weight, mesh, rowwise_placement)) sharded_module = distribute_module(MyModule(), mesh, partition_fn=shard_fc)
2.3 混合並行
混合並行（Hybrid Parallelism，HP）是將多種並行策略如資料並行、流水線並行和張量並行等進行混合使用。透過結合不同的並行策略，混合並行可以充分發揮各種並行策略的優點，以最大程度地提高計算效能和效率。
針對千億規模的大語言模型，通常在每個伺服器內部使用張量並行策略，由於該策略涉及的網路通訊量較大，需要利用伺服器內部的不同計算裝置之間進行高速通訊頻寬。透過流水線並行，將模型的不同層劃分為多個階段，每個階段由不同的機器負責計算。這樣可以充分利用多臺機器的計算能力，並透過機器之間的高速通訊來傳遞計算結果和中間資料，以提高整體的計算速度和效率。
最後，在外層疊加資料並行策略，以增加併發數量，提升整體訓練速度。透過資料並行，將訓練資料分發到多組伺服器上進行並行處理，每組伺服器處理不同的資料批次。這樣可以充分利用多臺伺服器的計算資源，並增加訓練的併發度，從而加快整體訓練速度。
BLOOM 使用了Megatron-DeepSpeed[104] 框架進行訓練，主要包含兩個部分：Megatron-LM 提供張量並行能力和資料載入原語；DeepSpeed提供ZeRO 最佳化器、模型流水線以及常規的分散式訓練元件。透過這種方式可以實現資料、張量和流水線三維並行，BLOOM 模型訓練時採用的平行計算結構如圖14所示。
BLOOM 模型訓練使用了由48 個NVIDIA DGX-A100 伺服器組成的叢集，每個DGX-A100 伺服器包含8 張NVIDIA A100 80GB GPU，總計包含384 張。BLOOM 訓練採用的策略是首先將叢集分為48 個一組，進行資料並行。
接下來，模型整體被分為12 個階段，進行流水線並行。每個階段的模型被劃分到4 張GPU 中，進行張量並行。同時BLOOM 也使用了ZeRO（零冗餘最佳化器）[134] 進一步降低了模型對視訊記憶體的佔用。用了透過上述四個步驟可以實現數百個GPU 的高效平行計算。
圖14 BLOOM 模型訓練時採用的平行計算結構
2.4 計算裝置記憶體最佳化
當前大語言模型訓練通常採用Adam 最佳化演算法，除了需要每個引數梯度之外，還需要一階動量（Momentum）和二階動量（Variance）。雖然Adam 最佳化演算法相較SGD 演算法通常效果更好也更穩定，但是對計算裝置記憶體的佔用顯著增大。
為了降低記憶體佔用，大多數系統已經採用了混合精度訓練（Mixed Precision Training）方式，即同時存在FP16（16 位浮點數）或者BF16（Bfloat16）和FP32（32 位浮點數）兩種格式的數值。FP32、FP16 和BF16 表示如圖4.15所示。FP32 中第31 位為符號位，第30 到第23 位用於表示指數，第22 到第0 位用於表示尾數。FP16 中第15 位為符號位，第14 到第10 位用於表示指數，第9 到第用於表示尾數。BF16 中第15 位為符號位，第14 到第7 位用於表示指數，第6 到第0 位用於表示尾數。由於FP16 的值區間比FP32 的值區間小很多，所以在計算過程中很容易出現上溢位和下溢位。BF16 相較於FP16 以精度換取更大的值區間範圍。但是，由於FP16 和BF16 相較FP32 精度低，訓練過程中可能會出現梯度消失和模型不穩定的問題。
因此，需要使用一些技術來解決這些問題，例如動態損失縮放（Dynamic Loss Scaling）和混合精度最佳化器（Mixed Precision Optimizer）等。混合精度最佳化的過程如圖4.16所示。Adam 最佳化器狀態包括採用FP32 儲存的模型引數備份，一階動量和二階動量也都採用FP32 格式儲存。假設模型引數量為Φ，模型引數和梯度都是用FP16格式儲存，則共需要2Φ + 2Φ + (4Φ + 4Φ + 4Φ) = 16Φ 位元組儲存。
其中Adam 狀態佔比75%。動態損失縮放反向傳播前，將損失變化（dLoss）手動增大2K 倍，因此反向傳播時得到的啟用函式梯度則不會溢位；反向傳播後，將權重梯度縮小2K 倍，恢復正常值。舉例來說，對於包含75 億個引數模型，如果用FP16 格式，只需要15GB 計算裝置記憶體，但是在訓練階段模型狀態實際上需要耗費120GB。
計算卡記憶體佔用中除了模型狀態之外，還有剩餘狀態（Residual States），包括啟用值（Activation）、各種臨時緩衝區（Buffer）以及無法使用的視訊記憶體碎片（Fragmentation）等。由於啟用值可以用檢查點（Activation Checkpointing）方式使得啟用值記憶體佔用大幅度減少，因此如何減少模型狀態尤其是Adam 最佳化器狀態是解決記憶體佔用問題的關鍵。
圖16 混合精度最佳化過程
以上是我簡單介紹分散式機器學習系統的基礎概念、分散式訓練叢集架構、分散式訓練並行策略，DeepSpeed 為例如何在叢集上訓練大語言模型在下一次的文章給大家繼續介紹，歡迎大家點贊關注支援，您的支援是我創作的動力。
參考內容：
(1) 收藏丨30個大語言模型訓練相關的資料集分享 - 知乎. https://zhuanlan.zhihu.com/p/612243919.
(2) 大語言模型訓練資料常見的4種處理方法 - 知乎. https://zhuanlan.zhihu.com/p/673045395.
(3)《大規模語言模型：從理論到實踐》張奇等著. —北京：電子工業出版社
(4) 大語言模型綜述 - Renmin University of China. http://ai.ruc.edu.cn/research/science/20230605100.html.
點選關注，第一時間瞭解華為雲新鮮技術~