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

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

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

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

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

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

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

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

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

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

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

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

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

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

針對分布式訓(xùn)練服務(wù)器集群進行架構(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ù)器：提供大量的計算資源
• 參數(shù)服務(wù)器：提供充足的內(nèi)存資源和通信資源

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

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

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

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

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

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

常見通信原語如下：

(1) Broadcast

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

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

(2) Scatter

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

(3) Reduce

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

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

(4) All Reduce

在所有節(jié)點上執(zhí)行同樣的Reduce操作，如求和、求最小值、求最大值等。可通過單節(jié)點上Reduce+Broadcast操作完成。

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

(5) Gather

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

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

「All Gather」在所有節(jié)點上收集所有其他節(jié)點的數(shù)據(jù)，最終使每個節(jié)點都擁有一份完整的數(shù)據(jù)集合。可以視為Gather操作與Broadcast操作的結(jié)合體。如下圖所示，All Gather操作將所有計算設(shè)備上的數(shù)據(jù)收集到各個計算設(shè)備。

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

(7) All to All

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

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

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

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

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

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

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

前向計算（Forward Pass）：

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

反向計算（Backward Pass）：

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

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

• 對數(shù)據(jù)進行切分（Partition），并將同一個模型copy到多個設(shè)備上，每個設(shè)備并行執(zhí)行不同的數(shù)據(jù)分片，即「數(shù)據(jù)并行（Data Parallelism，DP）」。
• 對模型進行拆分，將模型中的算子分發(fā)到多個設(shè)備分別完成，即「模型并行（Model Parallelism，MP）」。
• 訓(xùn)練超大規(guī)模語言模型時，同時對數(shù)據(jù)和模型進行并行，即「混合并行（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核心思想是將一個大的batch數(shù)據(jù)分割成多個子batch，并將子batch分配給不同的GPU進行并行計算。如下圖將訓(xùn)練過程分為前向傳播和反向傳播詳細(xì)分析：

前向傳播：

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

反向傳播：

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

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

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

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

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

(2) DistributedDataParallel（DDP)

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

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

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

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

# 2）從命令行接收local_rank參數(shù)，該參數(shù)表示當(dāng)前GPU在本地機器上的編號，用于后續(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)前進程使用的GPU設(shè)備，創(chuàng)建對應(yīng)的device對象
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)

-----------------------------------------------------------------------------
# 模型設(shè)置
# 將模型封裝進DistributedDataParallel，
# 指定模型運行在local_rank對應(yīng)的GPU上，同時將模型移動到相應(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)境中對數(shù)據(jù)集進行采樣，確保每個進程處理不同的數(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)
-------------------------------------------------------------------------------
# 運行的時候需要設(shè)置
for epoch in range(num_epochs):
    #  在每個epoch開始時，更新DistributedSampler的epoch，確保數(shù)據(jù)的隨機重排
     trainloader.sampler.set_epoch(epoch)
    #  遍歷數(shù)據(jù)集，前向傳播計算預(yù)測值，計算損失，執(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模式下存在的效率瓶頸和資源分配不均等問題，每個GPU節(jié)點都具備獨立的數(shù)據(jù)加載能力，無需等待主GPU的數(shù)據(jù)分發(fā)。并且，可執(zhí)行模型訓(xùn)練的每一環(huán)節(jié)，包括前向傳播、損失計算和反向傳播，實現(xiàn)了真正的并行計算。

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

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

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

2.模型并行MP

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

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

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

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

(1) 張量并行

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

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

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

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

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

① 嵌入式表示（Embedding）

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

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

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

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

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

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

② 矩陣乘法（MatMul）

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

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

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

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

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

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

③ 交叉熵Loss計算（CrossEntropyLoss）

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

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

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

具體的計算步驟如下圖所示。

(2) 流水線并行

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

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

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

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

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

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

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

① 樸素流水線并行

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

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

將其劃分到兩個GPU上：

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

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

② 樸素層并行的缺點：

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

③ GPipe流水線并行

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

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

由此可以看出，在每一時刻，僅有一個1個GPU工作，并且每個timesep花費的時間也比較長，因為GPU需要跑完整個minibatch的前向傳播。

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

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

④ PipeDream流水線并行

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

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

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

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

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

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

3.混合并行HP

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

(1) DP+PP

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

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

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

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

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

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

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

• DDP (Distributed DataParallel)
• ZeRO

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

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

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

• 使用Tensor Parallel（TP），因為僅靠PipelineParallel（PP）不足以容納大型層。
• ZeRO

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

具有快速的節(jié)點間連接時：

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

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

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