GPT-3模型為何難以復現？這也許是分布式AI框架的最優設計（2）

3.后向重計算

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

Checkpointing 的核心思想 是在前向網絡中標記少量的 Tensor （被 Checkpointing 的 Tensor ），前向計算就只會保留這些被標記的 Tensor， 其余的前向的 activation，會通過在反向傳播中根據 Checkpointing 的 Tensor 臨時重新計算一遍前向得到。這樣就使得大量的 activation 不需要一直保存到后向計算，有效減少了大量 Tensor 的生命周期，使得內存復用效率大幅提升。

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

兩種亞線性內存優化通過不同的方式達到了顯存優化：Checkpointing 是通過額外的計算開銷換顯存， CPU offload 通過額外的傳輸開銷換顯存。

Checkpointing 優化

上圖展示了兩層 Transformer Layer 在做 Checkpointing 之前和之后的計算圖對比， 其中重要的區別是前后向之間的連邊從很多條變成了兩條。不同框架實現Checkpointing的思路不同，Megatron 是自己重載了 torch.nn.Module ，實現了自己的 checkpointed_forward，相當于定制化了 Transformer Layer 的前后向執行邏輯；OneFlow 的 Checkpointing 就是上圖中的設計， 我們在整個計算圖中插入了重計算的子圖，并使得后向對前向的消費轉移到了對重計算子圖的消費。

重計算并不是單獨為流水并行設計的，并且之前大多使用在單卡或者數據并行場景下。但這個優化在流水并行下就非常關鍵，因為它使得前向不需要緩存所有的 activation，而只需要緩存非常少個數的（比如一層 Transformer Layer 只會緩存一個 ）、被 checkpoint 的特定 Tensor ，從而大大節省了流水并行下的顯存開銷。

4. 1F1B 策略

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

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

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

1F1B 策略下的 Pipeline 時間線

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

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

Megatron ：PyTorch 分布式訓練的極限、痛點在哪兒？

NVIDIA 基于 PyTorch 開發了 Megatron，本質上是一個專用于 GPT 的模型庫，所有的代碼都是 Python 腳本，NVIDIA 為 GPT 專門定制了分布式訓練所需的算子、 流水并行調度器、模型并行所需的通信原語等功能。可以說，NVIDIA 在使用 PyTorch 做分布式訓練上已經做到極致了。

在本章節，我們會簡單介紹一下 Megatron 是如何使用 PyTorch 的，當你也了解 Megatron 的設計以后，你就可以回答這個問題： PyTorch 做分布式訓練，真的好用嗎？

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

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

而流水并行，模型網絡分布在各個設備上是非對稱的，各個設備“接力”執行網絡的一部分，這種并行方式用 PyTorch 要如何實現呢？

流水并行 2 卡接力執行網絡

上圖展示了流水并行下，前兩個 stage 分布在 GPU 0 和 GPU 1 上時，網絡的拓撲關系。GPU 0 和 GPU 1 是接力執行的， GPU 0 上的 T2 Layer 的輸出 Tensor 需要發給 GPU 1 上的 T3 Layer 作為輸入。

首先，你需要根據 stage 階段的不同，分別在各個設備上定義只屬于自己那部分的模型網絡，而由于第一個 stage 和最后一個 stage 在執行時序上的特殊性，這里 Megatron 還需要進行特判 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]

在每個設備根據自己的那部分網絡啟動以后， Megatron 需要給每個設備上的每一次執行前后都調用 NCCL 的通信操作，前一個 stage 的輸出需要通過 NCCL p2p的 ncclSend 操作發給 下一個 stage， 下一個 stage 必須同時調用 ncclRecv 進行接收。當這兩個操作成對出現時，這次傳輸才會成功。（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)

因此對于 PyTorch 用戶而言，用戶自己需要關心每個 stage 在什么時機需要 recv，什么時機要 send， 發給誰；同時根據 Pipeline 的執行時序，需要特判在前多少個 step，都是需要只做前向（因為后向還沒來）， 但又有一些 step，我需要既做前向又做后向，因此你可以看到在 megatron/p2p_communication.py 里，你會發現 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."""

通過這些接口，你就會發現，算法工程師如果想用 PyTorch 做流水并行，他需要精細的控制所有的流水細節， 包括每個 stage 的每個時刻是只做前向，還是前向后向一起做， 同時還需要管理不同 stage 之間收/發數據的節奏，這個要求對于用戶而言就太高了。

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

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

GPT 的大規模訓練需要同時用到數據并行、模型并行和流水并行， 對于一個邏輯上的 Transformer Layer，需要同時對一個層做數據并行和模型并行，這個在 Megatron 和 DeepSpeed 的語義里稱之為 data-parallel-size 和 tensor-model-parallel-size 。

為什么要既做數據并行，又做模型并行？其實是為了節省顯存，并充分利用 GPU 之間的高速互聯（NVLink 和 NVSwitch）帶寬與機器之間的 IB 網絡帶寬的差別，NVIDIA 設計了一種在機器間做數據并行， 在機器內做模型并行的混合并行。

在什么樣的網絡結構、參數規模、網絡拓撲下該用數據并行、模型并行還是流水并行，是一個非常復雜的問題。不同的并行方式導致的設備之間、機器之間的通信量是不同的；同時又需要考慮設備顯存的約束、 GPU 通信帶寬和網絡通信帶寬的占比、 總的 Batch Size 大小對模型收斂速度的影響等等。目前還沒有一個嚴格的理論來指導具體模型在具體網絡拓撲下究竟該用哪種并行配置最優。對于并行策略的研究，我們會在未來專門出一篇文章來探討這個話題。

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

于 GPT-3 必須使用模型并行，因此被放在了機器內做。

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

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

NVIDIA 模型并行通信推導

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

NVIDIA 在論文中實驗了相同的總模型并行度（ model-parallel-size = tensor-model-parallel-size * pipeline-model-parallel-size）下， 分配不同的模型并行和流水并行的 size，得出當 tensor-model-parallel-size = 8 時， 總的效率最高，這與每臺機器內的卡數相同 。

模型并行度和流水并行度對性能的影響

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

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)

這是最關鍵的一處數據同步操作。但即使在 GPT 這樣全部由 Transformer Layer 組成的非常規整的網絡里，模型并行需要插入的同步操作就包括但不限于：

AllReduce ：train_step 、 calc_params_l2_norm、CrossEntropy

Scatter ：RowParallelLinear

AllGather ：ColumnParallelLinear

等等...

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

哪處需要插入通信操作？（現在 GPT 的腳本里 NVIDIA 給推導了需要插入通信的位置，如果用戶想改網絡結構，想加/換一個Op，推導是不是都得重來？）

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

通信操作要跟誰通信？（數據并行和模型并行同時做時， 整個 GPU 集群會被分組，每一組組內做 AllReduce 同步數據， 組間在模型更新時 才同步模型梯度，這意味著每個 rank 的 GPU 想要通信時，是需要跟其他特定對應的 rank 做通信的，這更加增加了實現難度）

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

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

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