隨著大模型時(shí)代的到來，模型參數(shù)量、訓(xùn)練數(shù)據(jù)量、計(jì)算量等各方面急劇增長。大模型訓(xùn)練面臨新的挑戰(zhàn)：

• 顯存挑戰(zhàn)：例如，175B的GPT-3模型需要175B*4bytes即700GB模型參數(shù)空間，而常見的GPU顯存如A100是80G顯存，這樣看來連模型加載都困難更別說訓(xùn)練。
• 計(jì)算挑戰(zhàn)：175B的GPT-3模型計(jì)算量也很龐大了，再疊加預(yù)訓(xùn)練數(shù)據(jù)量，所需的計(jì)算量與BERT時(shí)代完全不可同日而語。

分布式訓(xùn)練(Distributed Training)則可以解決海量計(jì)算和內(nèi)存資源要求的問題。它可將一個(gè)模型訓(xùn)練任務(wù)拆分為多個(gè)子任務(wù)，并將子任務(wù)分發(fā)給多個(gè)計(jì)算設(shè)備（eg:單機(jī)多卡，多機(jī)多卡），從而解決資源瓶頸。

本文將詳細(xì)介紹分布式訓(xùn)練的基本概念、集群架構(gòu)、并行策略等，以及如何在集群上訓(xùn)練大語言模型。

一、何為分布式訓(xùn)練？

分布式訓(xùn)練是指將機(jī)器學(xué)習(xí)或深度學(xué)習(xí)模型訓(xùn)練任務(wù)分解成多個(gè)子任務(wù)，并在多個(gè)計(jì)算設(shè)備上并行訓(xùn)練，可以更快速地完成整體計(jì)算，并最終實(shí)現(xiàn)對(duì)整個(gè)計(jì)算過程的加速。

如上圖是單個(gè)計(jì)算設(shè)備和多個(gè)計(jì)算設(shè)備的不同，這里計(jì)算設(shè)備可以是CPU、GPU、TPU、NPU等。

在分布式訓(xùn)練的背景下，無論是單服務(wù)器內(nèi)的多計(jì)算設(shè)備還是跨服務(wù)器的多設(shè)備，系統(tǒng)架構(gòu)均被視為「分布式系統(tǒng)」。這是因?yàn)?，即使在同一服?wù)器內(nèi)部，多個(gè)計(jì)算設(shè)備(如GPU)之間的內(nèi)存也不一定是共享的，意味著「設(shè)備間的數(shù)據(jù)交換和同步必須通過網(wǎng)絡(luò)或高速互聯(lián)實(shí)現(xiàn)」，與跨服務(wù)器的設(shè)備通信本質(zhì)相同。

二、分布式訓(xùn)練集群架構(gòu)

分布式訓(xùn)練集群屬于高性能計(jì)算集群（High Performance Computing Cluster，HPC），其目標(biāo)是提供海量的計(jì)算能力。 在由高速網(wǎng)絡(luò)組成的高性能計(jì)算上構(gòu)建分布式訓(xùn)練系統(tǒng)。

高性能計(jì)算集群硬件組成如圖所示。

整個(gè)計(jì)算集群包含大量帶有計(jì)算加速設(shè)備的服務(wù)器，多個(gè)服務(wù)器會(huì)被安置在機(jī)柜中，服務(wù)器通過架頂交換機(jī)（Top of Rack Switch，ToR）連接網(wǎng)絡(luò)。在架頂交換機(jī)滿載的情況下，可以通過在架頂交換機(jī)間增加骨干交換機(jī)進(jìn)一步接入新的機(jī)柜。

每個(gè)服務(wù)器中通常是由2-16個(gè)計(jì)算加速設(shè)備組成，這些計(jì)算加速設(shè)備之間的高速通信直接影響到分布式訓(xùn)練的效率。傳統(tǒng)的PCI Express（PCIe）總線，即使是PCIe 5.0版本，也只能提供相對(duì)較低的128GB/s帶寬，這在處理大規(guī)模數(shù)據(jù)集和復(fù)雜模型時(shí)可能成為瓶頸。

為了解決這一問題，NVIDIA推出了NVLink和NVSwitch技術(shù)。如下圖所示，每個(gè)H100 GPU都有多個(gè)NVLink端口，并連接到所有四個(gè)NVSwitch上。每個(gè)NVSwitch都是一個(gè)完全無阻塞的交換機(jī)，完全連接所有8個(gè)H100計(jì)算加速卡。NVSwitch的這種完全連接的拓?fù)浣Y(jié)構(gòu)，使得服務(wù)器內(nèi)任何H100加速卡之間都可以達(dá)到900GB/s雙向通信速度。

針對(duì)分布式訓(xùn)練服務(wù)器集群進(jìn)行架構(gòu)涉及的主流架構(gòu)，目前主流的主要分為參數(shù)服務(wù)器（ParameterServer，簡稱PS）和去中心化架構(gòu)（Decentralized Network）兩種分布式架構(gòu)。

1.參數(shù)服務(wù)器架構(gòu)

參數(shù)服務(wù)器架構(gòu)的分布式訓(xùn)練系統(tǒng)中有兩種服務(wù)器：

• 訓(xùn)練服務(wù)器：提供大量的計(jì)算資源
• 參數(shù)服務(wù)器：提供充足的內(nèi)存資源和通信資源

如下所示是具有參數(shù)服務(wù)器的分布式訓(xùn)練集群的示意圖。在訓(xùn)練過程中，每個(gè)訓(xùn)練服務(wù)器都擁有完整的模型，并根據(jù)將分配到此服務(wù)器的訓(xùn)練數(shù)據(jù)集切片（Dataset Shard）進(jìn)行計(jì)算，將得到的梯度推送到相應(yīng)的參數(shù)服務(wù)器。參數(shù)服務(wù)器會(huì)等待兩個(gè)訓(xùn)練服務(wù)器都完成梯度推送，然后開始計(jì)算平均梯度，并更新參數(shù)。之后，參數(shù)服務(wù)器會(huì)通知訓(xùn)練服務(wù)器拉取最新的參數(shù)，并開始下一輪訓(xùn)練迭代。

2.去中心化架構(gòu)

去中心化架構(gòu)沒有中央服務(wù)器或控制節(jié)點(diǎn)，而是由節(jié)點(diǎn)之間進(jìn)行直接通信和協(xié)調(diào)，這種架構(gòu)的好處是可以減少通信瓶頸，提高系統(tǒng)的可擴(kuò)展性。

節(jié)點(diǎn)之間的分布式通信一般有兩大類：

• 點(diǎn)對(duì)點(diǎn)通信（Point-to-Point Communication）：在一組節(jié)點(diǎn)內(nèi)進(jìn)行通信
• 集合通信（Collective communication，CC）：在兩個(gè)節(jié)點(diǎn)之間進(jìn)行通信

去中心化架構(gòu)中通常采用集合通信實(shí)現(xiàn)。

常見通信原語如下：

(1) Broadcast

將數(shù)據(jù)從主節(jié)點(diǎn)發(fā)送到集群中的其他節(jié)點(diǎn)。如下圖，計(jì)算設(shè)備1將大小為1xN的張量廣播到其它設(shè)備，最終每張卡輸出均為1×N矩陣。

Broadcast在分布式訓(xùn)練中主要用于「模型參數(shù)初始化」。

(2) Scatter

主節(jié)點(diǎn)將一個(gè)大的數(shù)據(jù)塊分割成若干小部分，再將每部分分發(fā)到集群中的其他節(jié)點(diǎn)。如下圖，計(jì)算設(shè)備1將大小為1xN的張量分成4個(gè)子張量，再分別發(fā)送給其它設(shè)備。

(3) Reduce

將不同節(jié)點(diǎn)上的計(jì)算結(jié)果進(jìn)行聚合。Reduce操作可以細(xì)分為多種類型，包括SUM（求和）、MIN（求最小值）、MAX（求最大值）、PROD（乘積）、LOR（邏輯或）等，每種類型對(duì)應(yīng)一種特定的聚合方式。

如下圖所示，Reduce Sum操作將所有計(jì)算設(shè)備上的數(shù)據(jù)進(jìn)行求和，然后將結(jié)果返回到計(jì)算設(shè)備1。

(4) All Reduce

在所有節(jié)點(diǎn)上執(zhí)行同樣的Reduce操作，如求和、求最小值、求最大值等?？赏ㄟ^單節(jié)點(diǎn)上Reduce+Broadcast操作完成。

如下圖所示，All Reduce Sum操作將所有節(jié)點(diǎn)上的數(shù)據(jù)求和，然后將求和結(jié)果Broadcast到所有節(jié)點(diǎn)。

(5) Gather

將所有節(jié)點(diǎn)的數(shù)據(jù)收集到單個(gè)節(jié)點(diǎn)，可以看作是Scatter操作的逆操作。

如下圖所示，Gather操作將所有設(shè)備的數(shù)據(jù)收集到計(jì)算設(shè)備1中。

「All Gather」在所有節(jié)點(diǎn)上收集所有其他節(jié)點(diǎn)的數(shù)據(jù)，最終使每個(gè)節(jié)點(diǎn)都擁有一份完整的數(shù)據(jù)集合?？梢砸暈镚ather操作與Broadcast操作的結(jié)合體。如下圖所示，All Gather操作將所有計(jì)算設(shè)備上的數(shù)據(jù)收集到各個(gè)計(jì)算設(shè)備。

「Reduce Scatter」將每個(gè)節(jié)點(diǎn)的張量分割成多個(gè)塊，每個(gè)塊分發(fā)給不同的節(jié)點(diǎn)，再在每個(gè)節(jié)點(diǎn)執(zhí)行Reduce操作（如求和、平均等）。如下圖所示，Reduce Scatter操作將每個(gè)計(jì)算設(shè)備中的張量分割成4塊，并發(fā)送給4個(gè)不同的計(jì)算設(shè)備，每個(gè)計(jì)算設(shè)備對(duì)接收到的塊執(zhí)行Reduce Sum操作。

(7) All to All

將每個(gè)節(jié)點(diǎn)上的數(shù)據(jù)分割成多個(gè)塊，并將這些塊分別發(fā)送給不同的節(jié)點(diǎn)。

如下圖所示，All to All操作將每個(gè)計(jì)算設(shè)備中的張量分割成4塊，并發(fā)送給4個(gè)不同的計(jì)算設(shè)備。

三、分布式訓(xùn)練并行策略

分布式訓(xùn)練系統(tǒng)的核心目標(biāo)是將原本在單一計(jì)算節(jié)點(diǎn)上進(jìn)行的模型訓(xùn)練過程，轉(zhuǎn)化為能在多個(gè)計(jì)算節(jié)點(diǎn)上并行執(zhí)行，以加速訓(xùn)練速度并支持更大規(guī)模的模型和數(shù)據(jù)集。

在單節(jié)點(diǎn)模型訓(xùn)練中，系統(tǒng)結(jié)構(gòu)主要由兩大部分組成：數(shù)據(jù)和模型。訓(xùn)練過程由多個(gè)數(shù)據(jù)小批次（Mini-batch）完成。如圖所示，數(shù)據(jù)表示一個(gè)數(shù)據(jù)小批次。訓(xùn)練系統(tǒng)會(huì)利用數(shù)據(jù)小批次根據(jù)損失函數(shù)和優(yōu)化算法生成梯度，從而對(duì)模型參數(shù)進(jìn)行修正。

針對(duì)大語言模型多層神經(jīng)網(wǎng)絡(luò)的執(zhí)行過程，模型訓(xùn)練過程可以抽象為一個(gè)計(jì)算圖（Computational Graph）。這個(gè)圖由多個(gè)相互連接的算子（Operator）構(gòu)成，每個(gè)算子對(duì)應(yīng)神經(jīng)網(wǎng)絡(luò)中的一個(gè)層（Neural Network Layer），如卷積層、全連接層等。參數(shù)（Weights）則是這些層在訓(xùn)練過程中不斷更新的權(quán)重。

計(jì)算圖的執(zhí)行過程可以分為前向傳播和反向傳播兩個(gè)階段。

前向計(jì)算（Forward Pass）：

• 輸入數(shù)據(jù)：數(shù)據(jù)從輸入層開始，被送入計(jì)算圖的第一個(gè)算子。
• 算子執(zhí)行：每個(gè)算子接收輸入數(shù)據(jù)，執(zhí)行相應(yīng)的數(shù)學(xué)運(yùn)算（如矩陣乘法、激活函數(shù)等），并產(chǎn)生輸出。
• 數(shù)據(jù)傳遞：算子的輸出作為后續(xù)算子的輸入，沿著計(jì)算圖向前傳播。
• 輸出生成：當(dāng)數(shù)據(jù)到達(dá)計(jì)算圖的末端，即輸出層，產(chǎn)生最終的預(yù)測結(jié)果。

反向計(jì)算（Backward Pass）：

• 損失計(jì)算：在前向傳播完成后，使用損失函數(shù)比較預(yù)測輸出與實(shí)際標(biāo)簽，計(jì)算損失值。
• 梯度計(jì)算：從輸出層開始，反向遍歷計(jì)算圖，根據(jù)損失值和算子的導(dǎo)數(shù)，計(jì)算每個(gè)算子的梯度。
• 參數(shù)更新：利用計(jì)算出的梯度，根據(jù)選擇的優(yōu)化算法（如梯度下降、Adam等），更新模型參數(shù)。
• 傳播回溯：反向計(jì)算過程從輸出層向輸入層遞歸進(jìn)行，直到所有參數(shù)都被更新。

根據(jù)單設(shè)備模型訓(xùn)練流程，可以看出，如果進(jìn)行并行加速，可以從數(shù)據(jù)和模型兩個(gè)維度考慮：

• 對(duì)數(shù)據(jù)進(jìn)行切分（Partition），并將同一個(gè)模型copy到多個(gè)設(shè)備上，每個(gè)設(shè)備并行執(zhí)行不同的數(shù)據(jù)分片，即「數(shù)據(jù)并行（Data Parallelism，DP）」。
• 對(duì)模型進(jìn)行拆分，將模型中的算子分發(fā)到多個(gè)設(shè)備分別完成，即「模型并行（Model Parallelism，MP）」。
• 訓(xùn)練超大規(guī)模語言模型時(shí)，同時(shí)對(duì)數(shù)據(jù)和模型進(jìn)行并行，即「混合并行（Hybrid Parallelism，HP）」。

1.數(shù)據(jù)并行DP

數(shù)據(jù)并行是最常用的并行訓(xùn)練方式，主要分為DataParallel(DP)和DistributedDataParallel(DDP)兩種。

(1) DP

DP是早使用的數(shù)據(jù)并行方案，通過torch.nn.DataParallel()來調(diào)用，代碼如下:

# 設(shè)置可見的GPU
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1,2,3"

# 將模型放到GPU 0上，必須先把模型放在GPU上，后面才可以調(diào)用DP
model.cuda()

# 構(gòu)建DataParallel數(shù)據(jù)并行化
model=torch.nn.DataParallel(model)

DP核心思想是將一個(gè)大的batch數(shù)據(jù)分割成多個(gè)子batch，并將子batch分配給不同的GPU進(jìn)行并行計(jì)算。如下圖將訓(xùn)練過程分為前向傳播和反向傳播詳細(xì)分析：

前向傳播：

• 模型和完整的mini-batch數(shù)據(jù)被放置在Master GPU（例如GPU:0）上。
• GPU:0將mini-batch數(shù)據(jù)分割成若干個(gè)子batch，并將這些子batch分發(fā)（scatter）到其它GPU上。
• GPU:0將自身的模型參數(shù)復(fù)制到其它GPU，確保每個(gè)GPU上的模型參數(shù)完全相同。
• 每個(gè)GPU在單獨(dú)的線程上對(duì)其sub-mini-batch的數(shù)據(jù)前向傳播，計(jì)算出各自的輸出結(jié)果。
• GPU:0收集所有GPU的輸出結(jié)果。

反向傳播：

• GPU:0基于收集的輸出結(jié)果和真實(shí)label計(jì)算總損失loss，并得到loss的梯度。
• GPU:0將計(jì)算出的loss梯度分發(fā)（scatter）到所有GPU上。
• 每個(gè)GPU根據(jù)收到的loss梯度反向傳播，計(jì)算出所有模型參數(shù)的梯度。
• 所有GPU計(jì)算出的參數(shù)梯度被匯總回GPU:0。
• GPU:0基于匯總的梯度更新模型參數(shù)，完成一次迭代的訓(xùn)練。

有人說GPU:0好自私，把其它GPU當(dāng)做工具人，核心機(jī)密不對(duì)外，只給其他GPU數(shù)據(jù)，不給label，其它GPU得到結(jié)果它再給計(jì)算loss和loss梯度，然后分發(fā)給其他GPU去計(jì)算參數(shù)梯度，之后得到這些參數(shù)的梯度后再去更新參數(shù)，等下次需要其它GPU了再分發(fā)更新好的參數(shù)。

這是一個(gè)悲傷的故事，首先「單進(jìn)程多線程」就似乎已經(jīng)注定的結(jié)局，Python的全局解釋鎖給這些附屬的GPU戴上了沉沉的牢拷，其他GPU想奮起反抗，但是DP里面只有一個(gè)優(yōu)化器Optimizer，這個(gè)優(yōu)化器Optimizer只在Master GPU上進(jìn)行參數(shù)更新，當(dāng)環(huán)境不再不在改變的時(shí)候，其它GPU選擇了躺平，當(dāng)GPU:0忙前忙后去分發(fā)數(shù)據(jù)、匯總梯度，更新參數(shù)的時(shí)候，其它GPU就靜靜躺著。

DataParallel采用的是Parameter Server并行架構(gòu)，在實(shí)現(xiàn)多GPU或多節(jié)點(diǎn)并行訓(xùn)練時(shí)，存在一些固有的局限性：

• 通信開銷大：每個(gè)「計(jì)算節(jié)點(diǎn)」在每次迭代中都需要與參數(shù)服務(wù)器進(jìn)行多次通信，以獲取最新的參數(shù)更新并將計(jì)算的梯度發(fā)送回去。這種頻繁的通信會(huì)導(dǎo)致網(wǎng)絡(luò)帶寬成為瓶頸，尤其是當(dāng)模型參數(shù)量大且GPU數(shù)量眾多時(shí)，通信延遲和帶寬消耗會(huì)顯著影響整體訓(xùn)練速度。
• 負(fù)載不均衡：其中一個(gè)GPU被指定為Master GPU，負(fù)責(zé)匯總梯度和廣播更新等，Master GPU可能會(huì)承擔(dān)額外的通信和計(jì)算負(fù)擔(dān)，導(dǎo)致負(fù)載不均衡。這不僅會(huì)影響該GPU的計(jì)算效率，也可能拖慢整個(gè)訓(xùn)練過程的速度。同時(shí)導(dǎo)致GPU利用率不足。
• 僅支持單機(jī)多卡模式，無法實(shí)現(xiàn)多機(jī)多卡訓(xùn)練。

(2) DistributedDataParallel（DDP)

DDP采用多進(jìn)程架構(gòu)，賦予了每個(gè)GPU更多的自由，支持多機(jī)多卡分布式訓(xùn)練。每個(gè)進(jìn)程都擁有自己的優(yōu)化器Optimizer，可獨(dú)立優(yōu)化所有步驟。每個(gè)進(jìn)程中在自己的GPU上計(jì)算loss，反向計(jì)算梯度。

在DDP中，不存在所謂的Master GPU，所有GPU節(jié)點(diǎn)地位平等，共同參與訓(xùn)練過程的每一個(gè)環(huán)節(jié)。每個(gè)GPU都會(huì)計(jì)算loss和梯度，然后通過高效的通信協(xié)議（如AllReduce）與其它GPU同步梯度，確保模型參數(shù)的一致性。

實(shí)現(xiàn)代碼如下：

# 初始化分布式環(huán)境
-----------------------------------------------------------------------------
# 1) 指定通信后端為nccl（NVIDIA Collective Communications Library），
  # 這是針對(duì)GPU集群優(yōu)化的高性能通信庫
 torch.distributed.init_process_group(backend='nccl')

# 2）從命令行接收local_rank參數(shù)，該參數(shù)表示當(dāng)前GPU在本地機(jī)器上的編號(hào)，用于后續(xù)的設(shè)備設(shè)置
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
args = parser.parse_args()

# 3) 設(shè)置cuda
# 根據(jù)local_rank設(shè)置當(dāng)前進(jìn)程使用的GPU設(shè)備，創(chuàng)建對(duì)應(yīng)的device對(duì)象
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)

-----------------------------------------------------------------------------
# 模型設(shè)置
# 將模型封裝進(jìn)DistributedDataParallel，
# 指定模型運(yùn)行在local_rank對(duì)應(yīng)的GPU上，同時(shí)將模型移動(dòng)到相應(yīng)的設(shè)備
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
model.to(device)
-----------------------------------------------------------------------------
# Dataset設(shè)置
Test_data = FunDataset(args.input)
# 創(chuàng)建DistributedSampler，用于在分布式環(huán)境中對(duì)數(shù)據(jù)集進(jìn)行采樣，確保每個(gè)進(jìn)程處理不同的數(shù)據(jù)子集
test_sample = torch.utils.data.distributed.DistributedSampler(Test_data)
# 使用DataLoader加載數(shù)據(jù)，指定sampler為DistributedSampler，確保數(shù)據(jù)的分布式加載和處理
test_data_dataset = DataLoader(dataset=Test_data, batch_size=args.batch_size, shuffle=False,
                                   collate_fn=Test_data.collate__fn,
                                   drop_last=False,sampler=test_sample)  # , pin_memory=True)
-------------------------------------------------------------------------------
# 運(yùn)行的時(shí)候需要設(shè)置
for epoch in range(num_epochs):
    #  在每個(gè)epoch開始時(shí)，更新DistributedSampler的epoch，確保數(shù)據(jù)的隨機(jī)重排
     trainloader.sampler.set_epoch(epoch)
    #  遍歷數(shù)據(jù)集，前向傳播計(jì)算預(yù)測值，計(jì)算損失，執(zhí)行反向傳播和參數(shù)更新
     for data, label in trainloader:
        prediction = model(data)
        loss = loss_fn(prediction, label)
        loss.backward()
        optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
        optimizer.step()

DDP解決了DP模式下存在的效率瓶頸和資源分配不均等問題，每個(gè)GPU節(jié)點(diǎn)都具備獨(dú)立的數(shù)據(jù)加載能力，無需等待主GPU的數(shù)據(jù)分發(fā)。并且，可執(zhí)行模型訓(xùn)練的每一環(huán)節(jié)，包括前向傳播、損失計(jì)算和反向傳播，實(shí)現(xiàn)了真正的并行計(jì)算。

引入了DistributedSampler，用于在分布式環(huán)境中均勻分割數(shù)據(jù)集，確保每個(gè)GPU處理的數(shù)據(jù)互不重疊，避免了數(shù)據(jù)冗余和計(jì)算浪費(fèi)。

采用了高效的Ring All-Reduce算法作為通信后端，用于梯度的聚合和參數(shù)的同步。在每個(gè)訓(xùn)練迭代結(jié)束時(shí)，每個(gè)GPU計(jì)算出自己的梯度后，通過環(huán)形網(wǎng)絡(luò)結(jié)構(gòu)與其他GPU進(jìn)行梯度交換，最終每個(gè)GPU都能獲取到所有GPU的梯度信息。

針對(duì)DP來說，Dataloder的batch_size是針對(duì)所有卡訓(xùn)練batch_size的和，例如10卡，每張卡batch_size是20，那么就要指定batch_size為200。針對(duì)DDP來說，batch_size就是每張卡所訓(xùn)練使用的batch_size為20。

2.模型并行MP

模型并行（Model Parallelism）通常用于解決單節(jié)點(diǎn)內(nèi)存不足的問題。以GPT-3為例，該模型擁有1750億參數(shù)，如果每個(gè)參數(shù)都使用32位浮點(diǎn)數(shù)表示，那么模型需要占用700GB（即175G× 4 Bytes）內(nèi)存。即使使用16位浮點(diǎn)數(shù)表示，每個(gè)模型副本也需要350GB的內(nèi)存。單個(gè)加速卡（如NVIDIA H100）的顯存容量（80GB）顯然不足以容納整個(gè)模型。

模型并行從計(jì)算圖的切分角度，可以分為以下兩種：

• 按模型的「layer層切分」到不同設(shè)備，即「層間并行或算子間并行」（Inter-operator Parallelism），也稱之為「流水線并行」（Pipeline Parallelism，PP）。

• 將計(jì)算圖層內(nèi)的「參數(shù)切分」到不同設(shè)備，即「層內(nèi)并行或算子內(nèi)并行」（Intra-operator Parallelism），也稱之為「張量并行」（Tensor Parallelism，TP）。

(1) 張量并行

張量并行（Tensor Parallelism，TP）旨在通過將模型參數(shù)和中間張量在多個(gè)設(shè)備（如GPU）之間進(jìn)行切分，以克服單個(gè)設(shè)備內(nèi)存限制的問題。

張量并行要解決的兩個(gè)核心問題：如何合理地將參數(shù)切分到不同設(shè)備，以及如何確保切分后的計(jì)算結(jié)果在數(shù)學(xué)上保持一致。

大語言模型都是以Transformer結(jié)構(gòu)為基礎(chǔ)，Transformer結(jié)構(gòu)主要由以下三種算子構(gòu)成：

• 嵌入式表示（Embedding）
• 矩陣乘法（MatMul）
• 交叉熵?fù)p失（Cross Entropy Loss）

這三種類型的算子均需要設(shè)計(jì)對(duì)應(yīng)的張量并行策略，才可以將參數(shù)切分到不同設(shè)備。

① 嵌入式表示（Embedding）

對(duì)于Embedding算子，總詞表數(shù)很大，將導(dǎo)致單計(jì)算設(shè)備顯存無法容納Embedding層參數(shù)。

例如，詞表數(shù)為64000，嵌入表示維度為5120，使用32位浮點(diǎn)數(shù)表示，整層參數(shù)需要的顯存大約為6400x5120x4/1024/1024=1250MB，加上反向傳播時(shí)的梯度存儲(chǔ)，總計(jì)近2.5GB，這對(duì)于顯存有限的設(shè)備而言是一個(gè)嚴(yán)峻挑戰(zhàn)。

為了解決這一問題，可以采用張量并行策略，將Embedding層的參數(shù)按詞維度切分，每個(gè)計(jì)算設(shè)備只存儲(chǔ)部分詞向量，然后通過匯總各個(gè)設(shè)備上的部分詞向量，從而得到完整的詞向量。

如下圖所示是單節(jié)點(diǎn)Embedding和兩節(jié)點(diǎn)張量并行的示意圖。

在單節(jié)點(diǎn)上，執(zhí)行Embedding操作，bz是批次大?。╞atch size），Embedding的參數(shù)大小為[word_size, hidden_size]，計(jì)算得到[bz,hidden_size]張量。

在兩節(jié)點(diǎn)上，可以將Embedding層參數(shù)按詞維度切分為兩半，每臺(tái)設(shè)備存儲(chǔ)一半的詞向量，即參數(shù)大小為[word_size/2, hidden_size]，分別存儲(chǔ)在兩個(gè)設(shè)備上。在前向傳播過程中，每個(gè)設(shè)備根據(jù)輸入的詞匯ID查詢自身的詞向量。如果無法查到，則該詞的表示為0。各設(shè)備計(jì)算得到的詞向量結(jié)果形狀為[bz, hidden_size]，由于詞匯可能被分割在不同設(shè)備上，需要通過跨設(shè)備的All Reduce Sum操作，將所有設(shè)備上的詞向量結(jié)果進(jìn)行匯總求和，以得到完整的詞向量表示?？梢钥闯?，這里的輸出結(jié)果和單節(jié)點(diǎn)執(zhí)行的結(jié)果一致。

② 矩陣乘法（MatMul）

矩陣乘的張量模型并行充分利用矩陣分塊乘法的原理。

例如，要實(shí)現(xiàn)矩陣乘法Y=X*A。其中，X是維度為MxN的輸入矩陣，A是維度為NxK的參數(shù)矩陣，Y是輸出，維度為MxK。

當(dāng)參數(shù)矩陣A的尺寸過大，以至于單個(gè)卡無法容納時(shí)，可以將參數(shù)矩陣A切分到多張卡上，并通過集合通信匯集結(jié)果，保證最終結(jié)果在數(shù)學(xué)計(jì)算上等價(jià)于單卡計(jì)算結(jié)果。

這里A切分方式包括按列切塊和按行切塊兩種。

「按列切塊」。如下圖所示，將參數(shù)矩陣A按列方向切分為A1和A2。將子矩陣A1和A2分配到兩張卡上。在計(jì)算過程中，每張卡將執(zhí)行獨(dú)立的矩陣乘法，即卡1計(jì)算Y1=X*A1，卡2計(jì)算Y2=X*A2。計(jì)算完成后，通過All Gather操作，每張卡將獲取另一張卡上的計(jì)算結(jié)果。在收集到所有計(jì)算結(jié)果后，每張卡將拼接它們收到的片段，形成完整的Y矩陣。最終，無論在哪張卡上查看，Y都將是一個(gè)完整的MxK矩陣，與單卡計(jì)算的結(jié)果完全等價(jià)。

「按行切塊」。如下圖所示，將參數(shù)矩陣A按列方向切分為A1和A2。為了與按行切塊的A矩陣相乘，輸入矩陣X（尺寸為MxN）也需要按列方向切分為X1和X2。將子矩陣A1和A2分別分配到兩張卡上。同時(shí)，將X1和X2也分別分配到對(duì)應(yīng)的卡上。每張卡將執(zhí)行獨(dú)立的矩陣乘法，即卡1計(jì)算Y1=X1*A1，卡2計(jì)算Y2=X2*A2。計(jì)算完成后，通過All Reduce Sum操作，每張卡將匯總另一張卡上的計(jì)算結(jié)果。在收集到所有計(jì)算結(jié)果后，每張卡將整合它們收到的片段，形成完整的Y矩陣的行。同理，參數(shù)矩陣A按行切塊的張量模型并行策略，通過巧妙地切分矩陣和利用多卡的計(jì)算能力，有效地解決了單卡顯存限制的問題，同時(shí)確保了計(jì)算結(jié)果的數(shù)學(xué)等價(jià)性。

③ 交叉熵Loss計(jì)算（CrossEntropyLoss）

分類網(wǎng)絡(luò)中，最后一層通常使用softmax函數(shù)結(jié)合交叉熵?fù)p失（CrossEntropyLoss）來評(píng)估模型的預(yù)測結(jié)果與真實(shí)標(biāo)簽之間的差距。然而，當(dāng)類別數(shù)量非常大時(shí)，單個(gè)GPU卡可能無法存儲(chǔ)和計(jì)算logit矩陣，導(dǎo)致顯存溢出。為了解決這一問題，可以采用張量并行策略，將logit矩陣和標(biāo)簽按類別數(shù)維度切分，通過中間結(jié)果的通信，計(jì)算出全局的交叉熵?fù)p失。具體的計(jì)算步驟如下：

首先計(jì)算softmax值。其中，p表示張量并行的設(shè)備號(hào)。

得到Softmax結(jié)果之后，同時(shí)對(duì)標(biāo)簽Target按類別切分，每個(gè)設(shè)備得到部分損失，最后再進(jìn)行一次通信，得到所有類別的損失。

具體的計(jì)算步驟如下圖所示。

(2) 流水線并行

所謂流水線并行，就是由于模型太大，無法將整個(gè)模型放置到單張GPU卡中，因此，將模型的不同層放置到不同的計(jì)算設(shè)備，降低單個(gè)計(jì)算設(shè)備的顯存消耗，從而實(shí)現(xiàn)超大規(guī)模模型訓(xùn)練。

流水線并行PP（Pipeline Parallelism），是一種最常用的并行方式，將模型的各個(gè)層分段處理，并將每個(gè)段分布在不同的計(jì)算設(shè)備上，使得前后階段能夠流水式、分批進(jìn)行工作。

如下圖所示是一個(gè)包含四層的深度學(xué)習(xí)模型，被切分為三個(gè)部分，并分別部署在三個(gè)不同的計(jì)算設(shè)備（Device 0、Device 1 和 Device 2）上。

其中，第一層（Layer 1）放置在Device 0上。第二層（Layer 2）和第三層（Layer 3）放置在Device 1上。第四層（Layer 4）放置在Device 2上。

前向計(jì)算過程中，輸入數(shù)據(jù)首先在 Device 0 上通過Layer 1的計(jì)算得到中間結(jié)果，并將中間結(jié)果傳輸?shù)紻evice 1，再繼續(xù)在Device 1上計(jì)算得到Layer 2 和 Layer 3的輸出，并將模型Layer 3的輸出結(jié)果傳輸?shù)紻evice 2。在 Device 2 上，數(shù)據(jù)經(jīng)過Layer 4 的計(jì)算，得到最終的前向計(jì)算結(jié)果。反向傳播過程類似。

最后，各個(gè)設(shè)備上的網(wǎng)絡(luò)層會(huì)使用反向傳播過程計(jì)算得到的梯度更新參數(shù)。由于各個(gè)設(shè)備間傳輸?shù)膬H是相鄰設(shè)備間的輸出張量，而不是梯度信息，因此通信量較小。

通過將模型切分為多個(gè)部分并分布到不同的計(jì)算設(shè)備上，流水線并行策略有效地?cái)U(kuò)展了可用于訓(xùn)練的GPU顯存容量。這意味著原本無法在單一GPU上裝載的大模型，現(xiàn)在可以通過類似流水線的方式，利用更多GPU的顯存來承載訓(xùn)練中的模型參數(shù)、梯度、優(yōu)化器狀態(tài)以及激活值等數(shù)據(jù)，從而實(shí)現(xiàn)超大規(guī)模模型的高效訓(xùn)練。

① 樸素流水線并行

當(dāng)模型規(guī)模超過單個(gè)GPU的處理能力時(shí)，樸素層并行（Naive Layer Parallelism）是一種直觀的并行策略，將模型的不同層分配到不同的GPU上，實(shí)現(xiàn)模型的并行化訓(xùn)練。如下所示是一個(gè)4層序列模型：

output=L4(L3(L2(L1(input)))))

將其劃分到兩個(gè)GPU上：

• GPU1負(fù)責(zé)計(jì)算前兩層：intermediate=L2(L1(input))
• GPU2負(fù)責(zé)計(jì)算后兩層：output=L4(L3(intermediate))

整個(gè)樸素層并行前向傳播和后向傳播的過程如上圖所示。GPU1執(zhí)行Layer 1和Layer 2的前向傳播，并緩存激活（activations），再將Layer 2的輸出（intermediate）發(fā)送給GPU2。GPU2接收來自GPU1的中間結(jié)果，執(zhí)行Layer 3和Layer 4的前向傳播，然后計(jì)算損失后，開始反向傳播，計(jì)算Layer 4和Layer 3的梯度，并將Layer 3的梯度返回給GPU1。GPU1接收梯度，繼續(xù)完成Layer 2和Layer 1的反向傳播。

② 樸素層并行的缺點(diǎn)：

• 低GPU利用率：同一時(shí)刻，只有其中一個(gè)GPU在執(zhí)行計(jì)算，其余GPU處于空閑狀態(tài)(又稱氣泡bubble)，這導(dǎo)致了計(jì)算資源的浪費(fèi)。
• 計(jì)算和通信沒有重疊：在數(shù)據(jù)傳輸期間，無論是前向傳播的中間結(jié)果還是反向傳播的梯度，GPU都處于等待狀態(tài)，這進(jìn)一步降低了計(jì)算效率。
• 高顯存占用：GPU1需要保存整個(gè)mini-batch的所有激活，直到反向傳播完成。如果mini-batch很大，這將顯著增加顯存的需求，可能超出單個(gè)GPU的顯存容量。

③ GPipe流水線并行

GPipe通過引入微批次（Micro-batch）的概念，顯著提高了模型并行訓(xùn)練的效率。將一個(gè)大的mini-batch進(jìn)一步細(xì)分為多個(gè)更小的、相等大小的微批次（microbatches），并在流水線并行的框架下獨(dú)立地對(duì)每個(gè)microbatch執(zhí)行前向傳播和反向傳播。然后將每個(gè)mircobatch的梯度相加，就能得到整個(gè)batch的梯度。由于每個(gè)層僅在一個(gè)GPU上，對(duì)mircobatch的梯度求和僅需要在本地進(jìn)行即可，不需要通信。

假設(shè)我們有4個(gè)GPU，并將模型按層劃分為4個(gè)部分，每個(gè)部分部署在一個(gè)GPU上。樸素層并行的過程如下所示：

由此可以看出，在每一時(shí)刻，僅有一個(gè)1個(gè)GPU工作，并且每個(gè)timesep花費(fèi)的時(shí)間也比較長，因?yàn)镚PU需要跑完整個(gè)minibatch的前向傳播。

GPipe將minibatch劃分為4個(gè)microbatch，然后依次送入GPU0。GPU0前向傳播后，再將結(jié)果送入GPU1，以此類推。整個(gè)過程如下所示。GPU0的前向計(jì)算被拆解為F11、F12、F13、F14，在GPU0中計(jì)算完成F11后，會(huì)在GPU1中開始計(jì)算F21，同時(shí)GPU0并行計(jì)算F12。

相比于樸素層并行方法，GPipe流水線方法可以有效降低并行氣泡大小。但是GPipe只有當(dāng)一個(gè)Mini-batch(4個(gè)Microbatch)中所有的前向傳播計(jì)算完成后，才能開始執(zhí)行反向傳播計(jì)算。因此還是會(huì)產(chǎn)生很多并行氣泡，從而降低了系統(tǒng)的并行效率。每個(gè)GPU需要緩存4份中間激活值。

④ PipeDream流水線并行

PipeDream流水線并行采用1F1B策略，即一個(gè)前向通道和一個(gè)后向通道，采用任務(wù)調(diào)度機(jī)制，使得下游設(shè)備能夠在等待上游計(jì)算的同時(shí)執(zhí)行其他可并行的任務(wù)，從而提高設(shè)備的利用率。

1F1B策略分為非交錯(cuò)式和交錯(cuò)式兩種方式調(diào)度方式。如下圖所示。

1F1B非交錯(cuò)式調(diào)度分為三個(gè)階段：

• 熱身階段：在這個(gè)階段，計(jì)算設(shè)備執(zhí)行不同數(shù)量的前向計(jì)算，為后續(xù)階段做準(zhǔn)備。
• 前向-后向階段：設(shè)備按照順序執(zhí)行一次前向計(jì)算，緊接著進(jìn)行一次后向計(jì)算。這一階段循環(huán)執(zhí)行，直到所有microbatch被處理完畢。
• 后向階段：在完成所有前向計(jì)算后，設(shè)備執(zhí)行剩余的后向計(jì)算，直至所有計(jì)算任務(wù)完成。

1F1B交錯(cuò)式調(diào)度要求microbatch的數(shù)量是流水線階段的整數(shù)倍。 每個(gè)設(shè)備不再僅負(fù)責(zé)連續(xù)多個(gè)層的計(jì)算，而是可以處理多個(gè)層的子集，這些子集被稱為模型塊。例如：

• 在傳統(tǒng)模式下，設(shè)備1可能負(fù)責(zé)層1-4，設(shè)備2負(fù)責(zé)層5-8。
• 在交錯(cuò)式模式下，設(shè)備1可以處理層1、2、9、10，設(shè)備2處理層3、4、11、12，以此類推。 這樣，每個(gè)設(shè)備在流水線中被分配到多個(gè)階段，可以并行執(zhí)行不同階段的計(jì)算任務(wù)，從而更好地利用流水線并行的優(yōu)勢(shì)。

3.混合并行HP

在進(jìn)行上百億/千億級(jí)以上參數(shù)規(guī)模的超大模型預(yù)訓(xùn)練時(shí)，通常會(huì)組合上述(數(shù)據(jù)并行、張量并行、流水線并行)多種并行技術(shù)一起使用。常見的分布式并行技術(shù)組合方案如下。

(1) DP+PP

DP rank 0看不到GPU2， DP rank 1看不到GPU3，對(duì)于DP，只有GPU 0和1提供數(shù)據(jù)，就好像只有2個(gè)GPU一樣。GPU0使用PP將其部分負(fù)載卸載到GPU2，同樣的GPU1使用PP將其部分負(fù)載卸載到GPU3。

由于每個(gè)維度至少需要2個(gè)GPU，因此這里至少需要4個(gè)GPU。

(2) 3D并行(DP+PP+TP)

3D并行是由數(shù)據(jù)并行(DP)、張量并行(TP)和流水線并行(PP)組成。由于每個(gè)維度至少需要2個(gè) GPU，因此這里至少需要8個(gè)GPU。

四、分布式訓(xùn)練并行策略選擇

1.單節(jié)點(diǎn)并行化策略

當(dāng)單個(gè)GPU可以裝入整個(gè)模型時(shí)：

• DDP (Distributed DataParallel)
• ZeRO

當(dāng)單個(gè)GPU無法裝入整個(gè)模型時(shí)：

• Pipeline Parallel (PP)
• Tensor Parallel (TP)
• ZeRO

當(dāng)單個(gè)GPU無法裝入模型的最大層時(shí)：

• 使用Tensor Parallel（TP），因?yàn)閮H靠PipelineParallel（PP）不足以容納大型層。
• ZeRO

2.多節(jié)點(diǎn)并行化策略

具有快速的節(jié)點(diǎn)間連接時(shí)：

• ZeRO
• 組合使用PP、TP和DP

節(jié)點(diǎn)間連接速度較慢且GPU內(nèi)存不足時(shí)：

• 組合使用PP、TP、DP和ZeRO