GPT-3模型為何難以復(fù)現(xiàn)？這也許是分布式AI框架的最優(yōu)設(shè)計(jì)（2）

3.后向重計(jì)算

Checkpointing 是陳天奇在2016年發(fā)表的論文 Training Deep Nets with Sublinear Memory Cost 中提到的，也稱之為亞線性內(nèi)存優(yōu)化。亞線性內(nèi)存優(yōu)化有兩種思路，Checkpointing 和 CPU offload：

Checkpointing 的核心思想 是在前向網(wǎng)絡(luò)中標(biāo)記少量的 Tensor （被 Checkpointing 的 Tensor ），前向計(jì)算就只會(huì)保留這些被標(biāo)記的 Tensor， 其余的前向的 activation，會(huì)通過在反向傳播中根據(jù) Checkpointing 的 Tensor 臨時(shí)重新計(jì)算一遍前向得到。這樣就使得大量的 activation 不需要一直保存到后向計(jì)算，有效減少了大量 Tensor 的生命周期，使得內(nèi)存復(fù)用效率大幅提升。

CPU offload 的思路類比于計(jì)算機(jī)操作系統(tǒng)中的“虛擬內(nèi)存”技術(shù)（將不常用的內(nèi)存臨時(shí)換入換出到磁盤上，從而增加內(nèi)存總量），在深度學(xué)習(xí)中，GPU 顯存（Device Memory）的特點(diǎn)是昂貴、高速且容量小，而 CPU 主存（Host Memory）的特點(diǎn)是便宜、相對(duì)低速和大容量；那么將前向計(jì)算中的一些暫時(shí)用不到的 activation 臨時(shí)換出到 CPU 主存上，等到反向計(jì)算需要時(shí)再換入到 GPU 顯存里，通過這種方式也可以節(jié)省顯存。

兩種亞線性內(nèi)存優(yōu)化通過不同的方式達(dá)到了顯存優(yōu)化：Checkpointing 是通過額外的計(jì)算開銷換顯存， CPU offload 通過額外的傳輸開銷換顯存。

Checkpointing 優(yōu)化

上圖展示了兩層 Transformer Layer 在做 Checkpointing 之前和之后的計(jì)算圖對(duì)比， 其中重要的區(qū)別是前后向之間的連邊從很多條變成了兩條。不同框架實(shí)現(xiàn)Checkpointing的思路不同，Megatron 是自己重載了 torch.nn.Module ，實(shí)現(xiàn)了自己的 checkpointed_forward，相當(dāng)于定制化了 Transformer Layer 的前后向執(zhí)行邏輯；OneFlow 的 Checkpointing 就是上圖中的設(shè)計(jì)， 我們?cè)谡麄€(gè)計(jì)算圖中插入了重計(jì)算的子圖，并使得后向?qū)η跋虻南M(fèi)轉(zhuǎn)移到了對(duì)重計(jì)算子圖的消費(fèi)。

重計(jì)算并不是單獨(dú)為流水并行設(shè)計(jì)的，并且之前大多使用在單卡或者數(shù)據(jù)并行場景下。但這個(gè)優(yōu)化在流水并行下就非常關(guān)鍵，因?yàn)樗沟们跋虿恍枰彺嫠械?activation，而只需要緩存非常少個(gè)數(shù)的（比如一層 Transformer Layer 只會(huì)緩存一個(gè) ）、被 checkpoint 的特定 Tensor ，從而大大節(jié)省了流水并行下的顯存開銷。

4. 1F1B 策略

除了重計(jì)算，上述 GPipe 的流水并行策略還有另外一個(gè)內(nèi)存問題，就是需要緩存幾份 activation，是等于一個(gè) batch 里有多少個(gè) micro-batch 的（梯度累加的次數(shù)）。通常，這個(gè)累加次數(shù)都比較大（為了盡可能流水，累加次數(shù)一般大于兩倍的 stage 數(shù)），那么即使緩存少數(shù) Tensor， 這種策略仍需要較多顯存。

因此，在另一篇流水并行的論文PipeDream （2018） 里就提出了改進(jìn)方法，稱之為 1F1B （One Forward pass followed by One Backward pass）的策略。這種改進(jìn)策略可以解決緩存 activation 的份數(shù)問題，使得 activation 的緩存數(shù)量只跟 stage 數(shù)相關(guān)，從而進(jìn)一步節(jié)省顯存，訓(xùn)練更大的模型。

1F1B 策略的出發(fā)點(diǎn)也比較直觀：由于前向計(jì)算的 activation 需要等到對(duì)應(yīng)的后向計(jì)算完成后才能釋放（無論有沒有使用 Checkpointing 技術(shù)），因此在流水并行下，如果想盡可能節(jié)省緩存 activation 的份數(shù)，就要盡量縮短每份 activation 保存的時(shí)間，也就是讓每份 activation 都盡可能早的釋放，所以要讓每個(gè) micro-batch 的數(shù)據(jù)盡可能早的完成后向計(jì)算，因此需要把后向計(jì)算的優(yōu)先級(jí)提高，讓 micro-batch 標(biāo)號(hào)小的后向比 micro-batch 標(biāo)號(hào)大的前向先做。因此，如果我們讓最后一個(gè) stage 在做完一次 micro-batch 的前向后，立馬就做本 micro-batch 的后向，那么我們就能讓其他的 stage 盡可能早的開始后向計(jì)算，這就是 1F1B 策略。其時(shí)間線如下圖所示：

1F1B 策略下的 Pipeline 時(shí)間線

從上圖 1F1B 和之前 GPipe 的流水線對(duì)比可知， GPipe 需要緩存 8 份的 activation 供后向使用，而 1F1B 策略只需要緩存 4 份。二者雖然空閑時(shí)間的占比是一樣的，但節(jié)省顯存就可以跑更多的 Layer 層數(shù) 和 更大的 micro-batch size，從而提升性能。

以上幾個(gè)關(guān)鍵技術(shù)（GPipe、梯度累加、重計(jì)算和 1F1B）的介紹就是分布式訓(xùn)練 GPT 的流水并行的核心技術(shù)（數(shù)據(jù)&模型并行我們放在下一章節(jié)詳細(xì)介紹）。無論是 NVIDIA 的Megatron（PyTorch），還是 OneFlow、PaddlePaddle、MindSpore ，都是通過不同的設(shè)計(jì)實(shí)現(xiàn)了上述相同的功能，而且 Megatron 在 NVIDIA 的深度優(yōu)化下， 在 GPU 上的性能表現(xiàn)已經(jīng)非常優(yōu)異了。那么 OneFlow 再搞一套 GPT 的意義何在？別急，看了下一章節(jié)，你就知道 PyTorch 做到上述這些技術(shù)的痛點(diǎn)在哪兒了。

Megatron ：PyTorch 分布式訓(xùn)練的極限、痛點(diǎn)在哪兒？

NVIDIA 基于 PyTorch 開發(fā)了 Megatron，本質(zhì)上是一個(gè)專用于 GPT 的模型庫，所有的代碼都是 Python 腳本，NVIDIA 為 GPT 專門定制了分布式訓(xùn)練所需的算子、 流水并行調(diào)度器、模型并行所需的通信原語等功能?？梢哉f，NVIDIA 在使用 PyTorch 做分布式訓(xùn)練上已經(jīng)做到極致了。

在本章節(jié)，我們會(huì)簡單介紹一下 Megatron 是如何使用 PyTorch 的，當(dāng)你也了解 Megatron 的設(shè)計(jì)以后，你就可以回答這個(gè)問題： PyTorch 做分布式訓(xùn)練，真的好用嗎？

1.流水并行，PyTorch 需要人工排線和精細(xì)控制流水

PyTorch 是單卡視角，一個(gè)設(shè)備上的 Tensor、模型腳本跟另一個(gè)設(shè)備上的 Tensor、模型腳本并無直接關(guān)系，對(duì)于每個(gè)設(shè)備上的模型腳本都完全對(duì)稱的（Mirror）最簡單的數(shù)據(jù)并行來說，PyTorch 這樣的設(shè)計(jì)沒有什么明顯的缺陷。每個(gè)設(shè)備上的腳本運(yùn)行到相同 batch 的模型更新部分（Optimizer），統(tǒng)一做一次模型同步（AllReduce 操作）就完成了數(shù)據(jù)并行，這就是 PyTorch 的 DDP（DistributedDataParallel）模塊。

而流水并行，模型網(wǎng)絡(luò)分布在各個(gè)設(shè)備上是非對(duì)稱的，各個(gè)設(shè)備“接力”執(zhí)行網(wǎng)絡(luò)的一部分，這種并行方式用 PyTorch 要如何實(shí)現(xiàn)呢？

流水并行 2 卡接力執(zhí)行網(wǎng)絡(luò)

上圖展示了流水并行下，前兩個(gè) stage 分布在 GPU 0 和 GPU 1 上時(shí)，網(wǎng)絡(luò)的拓?fù)潢P(guān)系。GPU 0 和 GPU 1 是接力執(zhí)行的， GPU 0 上的 T2 Layer 的輸出 Tensor 需要發(fā)給 GPU 1 上的 T3 Layer 作為輸入。

首先，你需要根據(jù) stage 階段的不同，分別在各個(gè)設(shè)備上定義只屬于自己那部分的模型網(wǎng)絡(luò)，而由于第一個(gè) stage 和最后一個(gè) stage 在執(zhí)行時(shí)序上的特殊性，這里 Megatron 還需要進(jìn)行特判 megatron/training.py 。

def train_step(...):
    if mpu.is_pipeline_first_stage():
        unwrapped_model = model[0]
    elif mpu.is_pipeline_last_stage():
        unwrapped_model = model[-1]

在每個(gè)設(shè)備根據(jù)自己的那部分網(wǎng)絡(luò)啟動(dòng)以后， Megatron 需要給每個(gè)設(shè)備上的每一次執(zhí)行前后都調(diào)用 NCCL 的通信操作，前一個(gè) stage 的輸出需要通過 NCCL p2p的 ncclSend 操作發(fā)給 下一個(gè) stage， 下一個(gè) stage 必須同時(shí)調(diào)用 ncclRecv 進(jìn)行接收。當(dāng)這兩個(gè)操作成對(duì)出現(xiàn)時(shí)，這次傳輸才會(huì)成功。（megatron/schedules.py）

def forward_backward_pipelining_without_interleaving(...):
    for i in range(num_microbatches_remaining):
        output_tensor = forward_step(...)
        if forward_only:
            p2p_communication.send_forward(output_tensor, timers)
        else:
            output_tensor_grad = p2p_communication.send_forward_recv_backward(output_tensor, timers)
        # Add input_tensor and output_tensor to end of list, then pop from the
        # start of the list for backward pass.
        input_tensors.append(input_tensor)
        output_tensors.append(output_tensor)
        if forward_only:
            if not last_iteration:
                input_tensor = p2p_communication.recv_forward(timers)
        else:
            input_tensor, output_tensor = input_tensors.pop(0), output_tensors.pop(0)
            input_tensor_grad = backward_step(...)
            if last_iteration:
                input_tensor = None
                p2p_communication.send_backward(input_tensor_grad, timers)
            else:
                input_tensor = p2p_communication.send_backward_recv_forward(input_tensor_grad, timers)

因此對(duì)于 PyTorch 用戶而言，用戶自己需要關(guān)心每個(gè) stage 在什么時(shí)機(jī)需要 recv，什么時(shí)機(jī)要 send， 發(fā)給誰；同時(shí)根據(jù) Pipeline 的執(zhí)行時(shí)序，需要特判在前多少個(gè) step，都是需要只做前向（因?yàn)楹笙蜻€沒來）， 但又有一些 step，我需要既做前向又做后向，因此你可以看到在 megatron/p2p_communication.py 里，你會(huì)發(fā)現(xiàn) Megatron 向用戶提供了這些操作：

def recv_forward(...):
    """Receive tensor from previous rank in pipeline (forward receive)."""
def recv_backward(...):
    """Receive tensor from next rank in pipeline (backward receive)."""
def send_forward(...):
    """Send tensor to next rank in pipeline (forward send)."""
def send_backward(...):
    """Send tensor to previous rank in pipeline (backward send)."""
def send_forward_recv_backward(...):
    """Batched send and recv with next rank in pipeline."""
def send_backward_recv_forward(...):
    """Batched send and recv with previous rank in pipeline."""
def send_forward_recv_forward(...):
    """Batched recv from previous rank and send to next rank in pipeline."""
def send_backward_recv_backward(...):
    """Batched recv from next rank and send to previous rank in pipeline."""
def send_forward_backward_recv_forward_backward(...):
    """Batched send and recv with previous and next ranks in pipeline."""

通過這些接口，你就會(huì)發(fā)現(xiàn)，算法工程師如果想用 PyTorch 做流水并行，他需要精細(xì)的控制所有的流水細(xì)節(jié)， 包括每個(gè) stage 的每個(gè)時(shí)刻是只做前向，還是前向后向一起做， 同時(shí)還需要管理不同 stage 之間收/發(fā)數(shù)據(jù)的節(jié)奏，這個(gè)要求對(duì)于用戶而言就太高了。

更讓人頭痛的是，PyTorch 并沒有機(jī)制保證這些流水并行中的各個(gè)設(shè)備之間數(shù)據(jù)交互的正確性 ，所以用戶不僅可能寫的不高效， 還可能寫錯(cuò)，即使寫錯(cuò)了，PyTorch 也無從檢查。 這些都給用戶帶來了極大的使用門檻。因此，也只有 NVIDIA 、 微軟等大企業(yè)的分布式訓(xùn)練專家可以搞得定 PyTorch 做流水并行。

2.模型并行，PyTorch 需要用戶在 kernel 中手寫通信原語操作，需要用戶推導(dǎo)所有的通信位置

GPT 的大規(guī)模訓(xùn)練需要同時(shí)用到數(shù)據(jù)并行、模型并行和流水并行， 對(duì)于一個(gè)邏輯上的 Transformer Layer，需要同時(shí)對(duì)一個(gè)層做數(shù)據(jù)并行和模型并行，這個(gè)在 Megatron 和 DeepSpeed 的語義里稱之為 data-parallel-size 和 tensor-model-parallel-size 。

為什么要既做數(shù)據(jù)并行，又做模型并行？其實(shí)是為了節(jié)省顯存，并充分利用 GPU 之間的高速互聯(lián)（NVLink 和 NVSwitch）帶寬與機(jī)器之間的 IB 網(wǎng)絡(luò)帶寬的差別，NVIDIA 設(shè)計(jì)了一種在機(jī)器間做數(shù)據(jù)并行， 在機(jī)器內(nèi)做模型并行的混合并行。

在什么樣的網(wǎng)絡(luò)結(jié)構(gòu)、參數(shù)規(guī)模、網(wǎng)絡(luò)拓?fù)湎略撚脭?shù)據(jù)并行、模型并行還是流水并行，是一個(gè)非常復(fù)雜的問題。不同的并行方式導(dǎo)致的設(shè)備之間、機(jī)器之間的通信量是不同的；同時(shí)又需要考慮設(shè)備顯存的約束、 GPU 通信帶寬和網(wǎng)絡(luò)通信帶寬的占比、 總的 Batch Size 大小對(duì)模型收斂速度的影響等等。目前還沒有一個(gè)嚴(yán)格的理論來指導(dǎo)具體模型在具體網(wǎng)絡(luò)拓?fù)湎戮烤乖撚媚姆N并行配置最優(yōu)。對(duì)于并行策略的研究，我們會(huì)在未來專門出一篇文章來探討這個(gè)話題。

對(duì)于大部分情況而言，數(shù)據(jù)并行的效率一般是最高的，但在 GPT-3 這樣的網(wǎng)絡(luò)參數(shù)規(guī)模下，單個(gè) GPU 根本裝不下這么大的模型，所以必須要用到模型并行和 流水并行來降低每個(gè) GPU 上的顯存需求。又基于 NVLink 和 IB 網(wǎng)絡(luò)通信帶寬的差別，NVIDIA 設(shè)計(jì)了一種折中的的方案，對(duì)整個(gè)集群拓?fù)渥龇纸M，分為機(jī)器間和機(jī)器內(nèi)，機(jī)器間的網(wǎng)絡(luò)傳輸速度較慢，往往是分布式并行的瓶頸，所以適合做流水并行和數(shù)據(jù)并行；機(jī)器內(nèi)的 NVLink 延遲低、帶寬高，正好符合模型并行的要求，由于 GPT-3 必須使用模型并行，因此被放在了機(jī)器內(nèi)做。

于 GPT-3 必須使用模型并行，因此被放在了機(jī)器內(nèi)做。

數(shù)據(jù)并行是在反向的梯度更新時(shí)需要插入 AllReduce 操作，而模型更新在 Gradient Accumulation 里是一個(gè)低頻操作，多個(gè) micro-batch 只會(huì)做一次，所以數(shù)據(jù)并行在機(jī)器間做是比較合適的。

模型并行（Tensor Model Parallelism）， NVIDIA 推導(dǎo)了 Transformer Layer 里的 MLP 和 Self-Attention 操作，模型并行下需要在特定位置插入 AllReduce 來實(shí)現(xiàn)前向、后向的數(shù)據(jù)同步工作。由于模型并行需要在每個(gè) micro-batch 的前向、后向都需要做數(shù)據(jù)同步， 屬于高頻操作， 所以 模型并行 適合在機(jī)器內(nèi)做。

NVIDIA 模型并行通信推導(dǎo)

流水并行的優(yōu)勢(shì)是帶寬需求比其它并行方式低，僅需要在 stage 之間傳輸數(shù)據(jù)， 同時(shí)還不會(huì)阻塞整個(gè)網(wǎng)絡(luò)的計(jì)算，因此在機(jī)器間做流水并行比較合適；但流水并行必須通過把一個(gè) Batch 分割成若干 micro-batches 才能發(fā)揮優(yōu)勢(shì)， 同時(shí)它還需要額外的顯存來緩存 activation，在 batch 間還會(huì)留下氣泡。

NVIDIA 在論文中實(shí)驗(yàn)了相同的總模型并行度（ model-parallel-size = tensor-model-parallel-size * pipeline-model-parallel-size）下， 分配不同的模型并行和流水并行的 size，得出當(dāng) tensor-model-parallel-size = 8 時(shí)， 總的效率最高，這與每臺(tái)機(jī)器內(nèi)的卡數(shù)相同 。

模型并行度和流水并行度對(duì)性能的影響

用 PyTorch 做模型并行的痛點(diǎn)是什么？如果你去了解一下 Megatron 搭 GPT 的模型腳本就立馬清楚了，我們知道模型并行需要在前后向插入一些數(shù)據(jù)同步的操作，但是在哪里插入？NVIDIA 給出了最主要的同步操作推導(dǎo)結(jié)果：在 RowParallelLinear 里需要將這個(gè)同步寫在 kernel 的 forward 函數(shù)里：

class RowParallelLinear(torch.nn.Module):
    def forward(self, input_):
        output_parallel = F.linear(input_parallel, self.weight)
        # All-reduce across all the partitions.
        output_ = reduce_from_tensor_model_parallel_region(output_parallel)

這是最關(guān)鍵的一處數(shù)據(jù)同步操作。但即使在 GPT 這樣全部由 Transformer Layer 組成的非常規(guī)整的網(wǎng)絡(luò)里，模型并行需要插入的同步操作就包括但不限于：

AllReduce ：train_step 、 calc_params_l2_norm、CrossEntropy

Scatter ：RowParallelLinear

AllGather ：ColumnParallelLinear

等等...

那么問題來了，算法工程師怎么知道這么長的模型腳本里，到底：

哪處需要插入通信操作？（現(xiàn)在 GPT 的腳本里 NVIDIA 給推導(dǎo)了需要插入通信的位置，如果用戶想改網(wǎng)絡(luò)結(jié)構(gòu)，想加/換一個(gè)Op，推導(dǎo)是不是都得重來？）

該插入什么通信操作？（除了 AllReduce，集合通信還有 ReduceScatter、AllGather、Reduce、Broadcast、All2All 等操作，除了集合通信，還有 Scatter 、Gather 等非對(duì)稱的切分、拼接操作，切分/拼接還要考慮對(duì) Tensor 的哪個(gè)維度操作...）

通信操作要跟誰通信？（數(shù)據(jù)并行和模型并行同時(shí)做時(shí)， 整個(gè) GPU 集群會(huì)被分組，每一組組內(nèi)做 AllReduce 同步數(shù)據(jù)， 組間在模型更新時(shí) 才同步模型梯度，這意味著每個(gè) rank 的 GPU 想要通信時(shí)，是需要跟其他特定對(duì)應(yīng)的 rank 做通信的，這更加增加了實(shí)現(xiàn)難度）

更要命的是，如果插入了通信操作，怎么保證正確性？PyTorch 沒法保證。PyTorch 將所有的操作都交給了用戶， 即使用戶插入了一個(gè)錯(cuò)誤的通信原語（比如將該插入 AllGather 操作的位置插入了 AllReduce），PyTorch 也沒法檢查出來。

所以這就是為什么只有 NVIDIA 可以用得了 PyTorch 做 Megatron，普通用戶只能直接用 megatron/pretrain_gpt.py，想基于 Megatron 做其他模型/網(wǎng)絡(luò)的遷移、二次開發(fā)和研究，是非常困難的。

其實(shí)，NVIDIA、 微軟、 PyTorch 都被繞進(jìn)一個(gè)大坑里去了：在沒有一致性視角（ Consistent View ）的情況下做復(fù)雜的分布式并行是非常困難的，往往只能做一些具體網(wǎng)絡(luò)具體場景具體算子的特判和分析，通過簡單的通信原語來實(shí)現(xiàn)分布式。而 OneFlow 通過一致性視角下的 Placement + SBP 就非常簡單的實(shí)現(xiàn)了通用的復(fù)雜并行支持。