為個人參與深度學習框架飛槳PaddlePaddle 開發時，梳理的個人筆記。
一、並行方式
1.資料並行（Batch維度）
資料並行分為了兩種模式：Data Parallel（DP） 和 Distributed Data Parallel（DDP） 。
1.1 Data Parallel
DP是一種單程序多執行緒的並行策略，只能在單機上進行訓練，從卡做Forward和Backward並行，主卡做梯度聚合和最佳化器更新，具體步驟如下：
- 單程序控制多GPU，即本質上是單程序多執行緒
- 首先將模型載入到主 GPU 上，再複製到各個指定從 GPU；
- 將輸入資料按照 Batch 維度進行拆分，各個 GPU 獨立進行 forward 計算；
- 將結果同步給主 GPU 完成梯度計算和引數更新，將更新後的權重引數複製到各個 GPU
存在的問題： 由於其是單程序控制多個GPU，故會存在GPU之間負載不均衡的問題，主GPU負載較大。
1.2 Distributed Data Parallel（DDP）
DDP 採用 AllReduce 架構，多程序的方式，突破鎖的束縛。在單機和多機上都可以使用。負載分散在每個 GPU 節點上，通訊成本（時間）是恆定的，與 GPU 數量無關，等於V/B（引數量/頻寬）。DDP不需要透過主GPU分發全模型的引數到每個GPU上。使用ring-all-reduce的方式進行通訊，隨著 GPU 數量 N 增加，總傳輸量恆定。也就是理論上，隨著GPU數量的增加，ring all-reduce有線性加速能力。
- 在飛槳中，
paddle.DataParallel介面預設提供的是DDP功能 - 提供了
no_sync()介面，用於暫停梯度同步的上下文管理器。在 no_sync()中引數梯度只會在模型上累加；直到 with 之外的第一個 forward-backward，梯度才會被同步。
>>> import numpy
>>> import paddle
>>> import paddle.distributed as dist
>>> from paddle.autograd import PyLayer
>>> from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... ctx.save_for_backward(y)
... return y
... @staticmethod
... def backward(ctx, dy):
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
>>> class SimpleNet(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.linear = paddle.nn.Linear(2, 2)
... def forward(self, inputs):
... inputs = cus_tanh.apply(inputs)
... return self.linear(inputs)
>>> if __name__ == '__main__':
... dist.init_parallel_env()
... model = SimpleNet()
... model = paddle.DataParallel(model)
... opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
... for step in range(10):
... x_data = numpy.random.randn(2, 2).astype(numpy.float32)
... x = paddle.to_tensor(x_data)
... x.stop_gradient = False
... # step 1 : skip gradient synchronization by 'no_sync'
... with model.no_sync():
... y_pred = model(x)
... loss = y_pred.mean()
... loss.backward()
... # step 2 : fuse + allreduce manually before optimization
... fused_allreduce_gradients(list(model.parameters()), None)
... opt.step()
... opt.clear_grad()
1.3 資料並行使用技巧
1.3.1 學習率設定
資料並行模式下學習率的設定技巧，其基本原則是學習率正比於 global batch size。 與單卡訓練相比，資料並行訓練通常有兩種配置：
- 一種是保持保持所有計算裝置的 batch size 的總和（我們稱為 global batch size）與單卡訓練的 batch size 保持一致。這種情形下，由於資料並行訓練和單卡訓練的 global batch size 是一致的，通常保持資料並行模式下各個計算裝置上的學習率與單卡訓練一致。
- 另一種情形是，保持資料並行模式下每個計算裝置的 batch size 和單卡訓練的 batch size 一致。這種情形下，資料並行模式的 global batch size 是單卡訓練的 N 倍。這裡， N 指的是資料平行計算的裝置數。因此，通常需要將資料並行模式下每個計算裝置的學習率相應的設定為單卡訓練的 N 倍。
- 這樣，資料並行模式下的初始學習率通常較大，不利於模型的收斂。因此，通常需要使用 warm-up 機制。即，在初始訓練時使用較小的學習率，並逐步緩慢增加學習率，經過一定迭代次數後，學習率增長到期望的學習率。
1.3.2 資料集切分
資料並行中，我們通常將資料集切分為 N 份，每個訓練卡負責訓練其中的一份資料。這裡， N 是資料並行的並行度。如我們前面介紹的，每一個迭代中，各個訓練卡均需要做一次梯度同步。因此，我們需要確保對於每個 epoch ，各個訓練卡經歷相同的迭代數，否則，執行迭代數多的訓練卡會一直等待通訊完成。實踐中，我們通常透過資料補齊或者丟棄的方式保證各個訓練卡經歷相同的迭代數。
- 資料補齊的方式指的是，為某些迭代數少訓練資料補充部分資料，從而保證切分後的各份資料集的迭代次數相同；
- 丟棄的方式則是丟棄部分迭代次數較多的資料，從而保證各份資料集的迭代次數相同。
通常，在每個 epoch 需要對資料做 shuffle 處理。因此，根據 shuffle 時機的不同，有兩種資料切分的方法。
- 一種是在資料切分前做 shuffle；即首先對完整的資料做 shuffle 處理，做相應的資料補充或丟棄，然後做資料的切分。
- 另一種是在資料切分後做 shuffle；即首先做資料的補充或丟棄和資料切分，然後對切分後的每一份資料分別做 shuffle 處理。
2.張量並行
總體而言，是將張量操作劃分到多個裝置上，以加速計算或增加模型大小；對模型每一層的層內引數進行切分，即對引數矩陣切片，並將不同切片放到不同GPU上；比如將原本在單卡中的矩陣乘法，切分到不同卡中進行矩陣乘法。訓練過程中，正向和反向傳播計算出的資料透過使用 All gather 或者 All reduce 的方法完成整合。
在Tansformer中，該策略會把 Masked Multi Self Attention 和 Feed Forward 都進行切分以並行化。利用 Transformers 網路的結構，透過新增一些同步原語來建立一個簡單的模型並行實現。張量並行適用於模型單層網路引數較大的情況。同時缺點也是十分明顯：
- 當環境是多機多卡，張量並行所需的all-reduce通訊需要跨伺服器進行連線，這比單機多GPU伺服器內的高頻寬通訊要慢（機間通訊比卡間通訊成本高）
- 高度的模型並行會產生很多小矩陣乘法，這可能會降低GPU的利用率。
張量模型並行需要解決兩個問題： 引數如何切分到不同裝置（切分方式）；以及切分後，如何保證數學一致性（數學等價）。本文以 NLP 中的 Transformer 結構為例，介紹張量模型並行的切分方式和隨機性控制
2.1 Embedding 切分
如下圖（a）所示。當採用模型並行時，Embedding 的引數被均勻切分到多個卡上。假設 Embedding 引數的維度為 N*D，並採用 K 張卡執行模型並行，那麼模型並行模式下每張卡上的 Embedding 引數的維度為 N//K*D 。當引數的維度 N 不能被卡數 K 整除時，最後一張卡的引數維度值為 (N//K+N%K)*D 。以下圖（b）為例，Embedding 引數的維度為 8*D ，採用 2 張卡執行模型並行，那麼每張卡上 Embedding 引數的維度為 4*D 。
為了便於說明，以下我們均假設 Embedding 的引數維度值 D 可以被模型並行的卡數 D 整除。此時，每張卡上 Embeeding 引數的索引值為 [0, N/K) ，邏輯索引值為 [k*N/K, (k+1)*N/K) ，其中 k 表示卡序號，0<=k<K。對於輸入索引 I，如果該索引在該卡表示的邏輯索引範圍內，則返回該索引所表示的表項（索引值為 I-k*N/K ；否則，返回值為全 0 的虛擬表項。隨後，透過 AllReduce 操作獲取所有輸出表項的和，即對應該 Embeding 操作的輸出；整個查表過程如下圖（b）所示。
2.2 Matmul 切分
2.2.1 列切分
對於矩陣乘操作，是按行或者列將矩陣切分 K 份。假設原始矩陣的維度為 M*N ，則按行切分後，各個卡上的矩陣維度為 M/K*N ；若按列切分，則各個卡上矩陣的維度值為 M*N/K 。圖（a）給出單卡上的矩陣乘法。圖（b）給出模型並行模式下的矩陣乘法，其中第二個矩陣按列切分到 2 張卡上；兩張卡分別得到結果矩陣的一部分。最後，透過 AllGather 通訊操作匯聚最終的結果。
2.2.2 行切分
下圖給出按行切分矩陣乘法的示例圖。其中，圖（a）給出單卡上的矩陣乘法。圖（b）給出模型並行模式下的矩陣乘法，其中第二個矩陣按行切分到 2 張卡上；第一個矩陣需要按列切分，以滿足矩陣乘法的維度要求；兩張卡分別得到結果矩陣的一部分。最後，透過 AllReduce 通訊操作按元素累加結果矩陣得到最終的結果。
相對於列切分，每張卡上的通訊量實際上是翻倍的？我們需要注意一下幾點：
- 模型並行下，需要確保模型並行組中各個卡讀取相同的資料；
- 模型並行下，除了被切分的運算元對應的輸出外，其它所有運算元的輸出在各個卡上是一致的。
3.流水線並行
通常來講，訓練更大規模的網路模型可以在多種任務上取得更好的效果，如提升影像分類任務的準確率。然而，隨著引數規模的擴大，AI 加速卡儲存（如 GPU 視訊記憶體）容量問題和卡的協同計算問題成為了訓練超大模型的瓶頸。流水線並行從模型切分和排程執行兩個角度解決了這些問題，下面將以飛槳流水線並行為例，介紹下基本原理和使用方法。
流水線原理是將不同的 layer 分配給指定 GPU 進行計算，流水線並行只需其之間點對點地通訊傳遞部分 activations。具體步驟包括：
- 在流水線並行之中，一個模型的各層會在多個GPU上做切分。
- 一個批次（batch）被分割成較小的微批（Micro-Batches），並在這些微批上進行流水線式執行。
- 透過流水線並行，一個模型的層被分散到多個裝置上。
- 當用於具有相同transformer塊重複的模型時，每個裝置可以被分配相同數量的transformer層。
- 在流水線模型並行中，訓練會在一個裝置上執行一組操作，然後將輸出傳遞到流水線中下一個裝置，下一個裝置將執行另一組不同操作。
流水線並行的方法，解決了超大模型無法在單裝置上裝下的難題，也解決了機器之間的通訊開銷的問題，使得每臺機器的資料傳輸量跟總的網路大小、機器總數、並行規模無關。如下圖，在最簡配置流水線並行模型下，任意時刻只有單個計算裝置處於計算狀態，其它計算裝置則處於空閒狀態，因此裝置利用率和計算效率較差。
為了最佳化流水線並行中裝置的計算效率，可以進一步將 mini-batch 切分成若干更小粒度的 micro-batch，以提升流水線並行的併發度，進而達到提升裝置利用率和計算效率的目的。如下圖所示，一個 mini-batch 被切分為 4 個 micro-batch；前向階段，每個裝置依次計算單個 micro-batch 的結果；從而增加了裝置間的併發度，降低了流水線並行 bubble 空間比例，提高了計算效率。
如上圖所示先進行前向計算，再進行反向計算，這種方式我們稱之為 F-the-B 模式。不難看出這種 F-then-B 模式由於快取了多個 micro-batch 的中間變數和梯度，視訊記憶體的實際利用率並不高。接下來我們介紹一種前向計算和反向計算交叉進行的方式，即 1F1B 模型。在 1F1B 模式下，前向計算和反向計算交叉進行，可以及時釋放不必要的中間變數。我們以下圖 1F1B 中 stage4 的 F42（stage4 的第 2 個 micro-batch 的前向計算）為例，F42 在計算前，F41 的反向 B41（stage4 的第 1 個 micro-batch 的反向計算）已經計算結束，即可釋放 F41 的中間變數，從而 F42 可以複用 F41 中間變數的視訊記憶體。1F1B 方式相比 F-then-B 方式峰值視訊記憶體可以節省 37.5%，對比樸素流水線並行峰值視訊記憶體明顯下降，裝置資源利用率顯著提升。
4.混合並行
5.MoE
通常來講，模型規模的擴充套件會導致訓練成本顯著增加，計算資源的限制成為了大規模密集模型訓練的瓶頸。為了解決這個問題， 《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》 提出了一種基於稀疏 MoE 層的深度學習模型架構，即將大模型拆分成多個小模型(專家, expert ), 每輪迭代根據樣本決定啟用一部分專家用於計算，達到了節省計算資源的效果； 並引入可訓練並確保稀疏性的門( gate )機制，以保證計算能力的最佳化。
與密集模型不同，MoE 將模型的某一層擴充套件為多個具有相同結構的專家網路( expert )，並由門( gate )網路決定啟用哪些 expert 用於計算，從而實現超大規模稀疏模型的訓練。 以上圖為例，示例模型包含 3 個模型層；如(a)到(b)，將中間層擴充套件為具有 n 個 expert 的 MoE 結構，並引入 Gating network 和 Top_k 機制，MoE 細節見圖(c)，計算過程如下述公式。
上述第 1 個公式表示了包含 n 個專家的 MoE 層的計算過程。具體來講，首先對樣本 x 進行門控計算， W 表示權重矩陣；然後由 Softmax 處理後獲得樣本 x 被分配到各個 expert 的權重； 然後只取前 k (通常取 1 或者 2）個最大權重，最終整個 MoE Layer 的計算結果就是選中的 k 個專家網路輸出的加權和。
import paddle
from paddle.nn import Layer, LayerList, Linear, Dropout
from paddle.incubate.distributed.models.moe import MoELayer
from paddle.distributed.collective import Group
from paddle.distributed import fleet
import numpy as np
# 構建一個可以正常訓練的模型
num_experts = 8
d_model = 512
d_hidden = 2048
class ExpertLayer(Layer):
def __init__(self, d_model, d_hidden, name=None):
super().__init__()
self.htoh4 = Linear(d_model, d_hidden)
self.h4toh = Linear(d_hidden, d_model)
def forward(self, x):
x = self.htoh4(x)
x = self.h4toh(x)
return x
# 然後初始化分散式環境，並構建 expert 通訊組 moe_group
fleet.init(is_collective=True)
moe_group = paddle.distributed.new_group(list(range(fleet.worker_num())))
# 設定門網路的 gate 策略和 top_k 機制，並將模型單層擴充套件為 num_expert 個相同結構的專家網路
gate_config = {
"type": "gshard",
"top_k": 2,
}
experts_list = LayerList()
for expi in range(num_experts):
exp_layer = ExpertLayer(d_model, d_hidden)
experts_list.append(exp_layer)
# 接著呼叫 MoELayer API 封裝並建立出 MoE 模型
class Model(Layer):
def __init__(self, d_model, d_hidden, name=None):
super().__init__()
self.linear1 = Linear(d_model, d_model)
self.moe_layer = MoELayer(d_model = d_model,
experts=experts_list,
gate=gate_config,
moe_group=moe_group,
recompute_interval=0)
self.linear2 = Linear(d_model, d_model)
self.dropout = Dropout(p=0.1)
def forward(self, x):
x = self.linear1(x)
x = self.moe_layer(x)
x = self.linear2(x)
x = self.dropout(x)
return x
model = Model(d_model, d_hidden)
optim = paddle.optimizer.SGD(parameters=model.parameters())
# 最後建立資料集，開始訓練
for step in range(1, 100):
x = paddle.rand([4, 256, d_model])
y = model(x)
loss = y.mean()
loss.backward()
optim.step()
optim.clear_grad()
print("=== step : {}, loss : {}".format(step, loss.numpy()))
# 執行方式：
# python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 --log_dir logs train_moe.py