本文分享自華爲雲社區《大模型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.