多 GPU 訓練策略如何選擇？

目錄

1. 簡介
2. 多 GPU 訓練策略概述
   1. 模型並行
   2. 張量並行
   3. 數據並行
   4. 流水線並行
   5. 序列並行
3. 模型並行策略
4. 張量並行策略
5. 數據並行策略
6. 流水線並行策略
7. 序列並行策略
8. 實際應用中的選擇
   1. 分散數據並行
   2. DDP Sharded
   3. DeepSpeed 策略
   4. FSDP
   5. DeepSpeed Stage 3
9. 結論
10. 參考資源

多GPU訓練策略概述 👥

在前一個視頻中，我們簡要介紹了多 GPU 訓練的概念。現在，讓我們更深入地討論不同的多 GPU 訓練策略。多 GPU 訓練可以分為幾個廣義的分類，包括模型並行、張量並行、數據並行、流水線並行和序列並行。讓我們逐一探討這些不同的並行策略以及它們的工作原理。

模型並行策略 🐇

模型並行策略通常在 GPU 記憶體有限的情況下非常有用。例如，如果我們有一個非常大的模型，我們可以將不同的層放在不同的 GPU 上以解決記憶體限制的問題。在 PyTorch 的代碼中，我們可以使用不同的 CUDA 調用將網絡的不同層放在不同的 GPU 上。然而，實際上我們不推薦使用這種策略，因為有更好的策略可以使用。

張量並行策略 🐘

張量並行策略與我們之前介紹的模型並行策略相關，但這裡的優勢在於我們對層進行水平拆分。也就是說，我們將層進行拆分，以便可以並行進行計算。在模型並行策略中，不同的層放在不同的 GPU 上，創建了一個順序瓶頸，這當然會降低速度。張量並行策略的有效性和效率高於模型並行策略。讓我們以矩陣乘法為例來說明這個概念。假設我們有一個2x3的矩陣和一個3x2的矩陣。如果右邊的矩陣有兩列，我們可以將其拆分為兩個可以並行進行計算的矩陣乘法問題。這樣，我們可以分別將左邊的矩陣與第一列向量和第二列向量進行乘法計算。然後，只要進行這兩個計算，我們就可以將這兩個結果連接起來，獲得乘法兩個矩陣時的原始結果。你可以看到，我們可以並行進行這兩個計算，這使得張量並行成為一種非常有效的多 GPU 訓練策略。當處理大型層和大型矩陣乘法時，這種策略非常有用。

數據並行策略 🐠

數據並行策略通常用於增加訓練吞吐量。在這種策略中，我們將批次分割為多個部分，並在不同的 GPU 上並行訓練模型。例如，我們可以創建模型的兩個副本，並將一個副本放在一個 GPU 上，將另一個副本放在另一個 GPU 上。然後，使用不同的數據點，我們可以並行訓練這兩個模型，然後在更新模型時平均梯度。這樣，我們可以通過在不同的 GPU 上並行訓練模型來更快地處理數據。值得注意的是，這只是一個簡要介紹，我們將在以後的視頻中詳細解釋這個概念。

流水線並行策略 ⛲

流水線並行策略是數據並行和模型並行的結合，其中模型被分割成不同的區塊。讓我們回顧一下最初的模型並行策略，其中不同的層在不同的 GPU 上，如下圖所示。F代表正向傳播，B代表反向傳播。這些不同的區塊從左到右執行。例如，如果第一層在GPU 4上，我們首先在GPU 4上執行計算。然後，我們去GPU 3上執行下一層，去GPU 2上執行下一層，最後去GPU 1上執行最後一層，然後以相反的順序進行反向傳播。流水線並行策略優化了 GPU 的並行工作。這意味著計算之間有更多的重疊，使得一個 GPU 不必等待另一個 GPU 完成。兩個 GPU 可以並行進行一些工作。

序列並行策略 🎺

序列並行策略是專門針對Transformer模型開發的，它將輸入序列分割成多個部分，以解決記憶體限制的問題。現在我們介紹了所有這些不同類型的多 GPU 訓練策略，那麼在實際應用中，我們應該使用哪種策略呢？個人而言，作為一個很好的默認選擇，我建議使用分散式數據並行策略（Distributed Data Parallel，DDP）並查看其效果。為了啟用這個策略，你只需要將strategy設置為DDP。如果你使用的是筆記本，請使用DDP_notebook。在某些情況下，您可能會遇到序列化的問題，這種情況下，您可能還需要使用DDP_spawn。但除非您遇到了pickling的問題，否則我建議使用DDP，因為它更快一些。請記住，在分散式數據並行中，我們將模型複製到不同的GPU上。如果我們有內存問題，如果模型對GPU內存太大，我們可能需要使用DDP_sharded。DDP_sharded是一種能夠使用額外的張量並行的數據並行策略。分割（sharding）這裡指的是對張量進行分割。作為DDP_sharded的替代方案，還有不同的DeepSpeed策略。例如，在DeepSpeed Stage 2中，它會對優化器和梯度進行分割。如果這還不足以解決內存限制問題，您還可以使用Stage 2的Offload策略。這裡的“offloading”意味著將優化器和梯度狀態轉移到CPU內存中。然而，將事物offloading到CPU內存會使事物變慢，因為在訓練過程中需要進行額外的傳輸。所以只有在由於內存限制而必須這樣做時，才建議使用offloading策略。現在，如果你只有少數一到幾個GPU，而訓練速度又不是很重要（例如，當你在較小的數據集上微調一個大型模型時），在僅有一個或少數幾個GPU時也有一些解決方法。對於這種情況，我建議在DeepSpeed中使用Stage 3策略。DeepSpeed Stage 3策略將優化器、梯度和權重參數都進行了分割。作為DeepSpeed Stage 3策略的替代方案，還有FSDP（Fully Sharded Data Parallel，完全分片的數據並行）。根據您的設備，其中一種可能會更快。實際上，我建議嘗試其中之一。在這個視頻中，我們簡要介紹了不同的並行策略，並為您提供了一些實際應用中使用的建議。在下一個視頻中，我們將更深入地研究數據並行策略及其在幕後的工作原理。

參考資源

[1] PyTorch官方文檔: https://pytorch.org/docs/stable/nn.html