在當今人工智能飛速發(fā)展的時代，大型語言模型（LLMs）正以其強大的語言理解和生成能力，改變著我們的生活和工作方式。在最近的一項研究中，科學家們?yōu)榱松钊肓私馊绾胃咝У赜柧毚笮驼Z言模型，進行了超過4000次的實驗。這些實驗動用了多達512個GPU（圖形處理單元），它們協(xié)同工作，為模型訓練提供了強大的計算支持。在這項研究中，研究人員特別關注了兩個關鍵指標：吞吐量（通過標記的大小來表示）和GPU利用率（通過標記的顏色來表示）。這兩個指標都根據模型的大小進行了標準化處理，以便更直觀地比較不同模型在不同硬件配置下的表現(xiàn)。

AI模型訓練的三大挑戰(zhàn)

內存使用：有限的“容器”

想象一下，你有一個裝滿水的杯子，但你還需要往里面加更多的水。如果杯子滿了，就再也裝不下任何東西了。這就像訓練AI模型時的內存問題。如果一個訓練步驟所需的內存超出了GPU的容量，那么訓練就無法繼續(xù)。內存是一個硬性限制，我們必須在有限的空間內完成復雜的計算任務。

計算效率：讓硬件“火力全開”

我們花了大價錢購買高性能的GPU，當然希望它們能時刻保持高效運轉。但現(xiàn)實往往是，GPU在等待數據傳輸或者等待其他GPU完成工作時，會浪費大量時間。這就像是一個生產線上的工人，因為原材料沒送到或者上一個環(huán)節(jié)還沒完成，只能干等著。為了提高效率，我們需要減少這些等待時間，讓硬件盡可能多地用于計算。

通信開銷：減少“內耗”

在多GPU協(xié)同工作的場景中，不同GPU之間需要頻繁地交換信息。但這種通信會占用大量的時間和資源，甚至會讓GPU處于閑置狀態(tài)。這就像是一個團隊開會，如果溝通成本過高，那么真正干活的時間就會減少。因此，我們需要巧妙地利用節(jié)點內部（速度快）和節(jié)點之間（速度慢）的帶寬，并盡可能讓通信和計算同時進行，從而減少通信開銷。

從單個GPU開始：AI模型訓練的第一步在探索如何用數千個GPU訓練大型AI模型之前，我們先從最基礎的部分開始——在單個GPU上訓練模型。這就好比在學會駕駛飛機之前，先要學會駕駛一輛汽車。別小看這一步，它是我們理解整個訓練過程的關鍵。

單GPU訓練的三個基本步驟

當你在單個GPU上訓練一個模型時，整個過程通?？梢苑譃槿齻€步驟：

前向傳播（Forward Pass）：將輸入數據通過模型，得到模型的輸出結果。這就好比把食材放進烤箱，等待美味的蛋糕出爐。

反向傳播（Backward Pass）：計算梯度，也就是找出模型需要改進的方向。這一步就像是檢查蛋糕的口感，看看哪里需要調整。

優(yōu)化步驟（Optimization）：用計算出的梯度來更新模型的參數，讓模型變得更好。這就好比根據反饋調整烤箱的溫度和時間。

關鍵參數：批次大?。˙atch Size）

批次大小（Batch Size）是訓練過程中一個非常重要的參數。它決定了每次訓練時輸入模型的數據量。批次大小的選擇對模型的訓練效果和效率有很大影響。

小批次大?。?/strong>在訓練初期，小批次大小可以幫助模型快速找到一個好的學習方向，就像在迷宮中快速試探出一條路。但如果一直使用小批次，模型的梯度會比較“嘈雜”，最終可能無法達到最優(yōu)性能。

• 大批量大?。?/strong>大批量大小可以提供更準確的梯度估計，但同時也會讓模型對每個訓練樣本的利用效率降低，導致訓練速度變慢，甚至浪費計算資源。

批次大小與訓練時間

批次大小還會影響訓練一個給定數據集所需的時間。小批次大小需要更多的優(yōu)化步驟來處理相同數量的數據，而每個優(yōu)化步驟都需要計算時間，因此總訓練時間會更長。不過，只要批次大小在最優(yōu)值附近，模型的最終性能通常不會受到太大影響。

在大型語言模型（LLM）的預訓練領域，批次大小通常用“token”（標記）數量而不是樣本數量來表示。這樣可以讓訓練數據與輸入序列長度無關，更加通用。

挑戰(zhàn)：內存不足怎么辦？

當我們嘗試將訓練擴展到更大的批次大小時，第一個挑戰(zhàn)就來了——GPU內存不足。當我們的GPU無法容納目標批次大小時，該怎么辦呢？

首先，我們需要理解為什么會出現(xiàn)內存不足的問題。在訓練神經網絡時，我們需要在內存中存儲以下內容：

模型權重（Weights）

模型梯度（Gradients）

優(yōu)化器狀態(tài)（Optimizer States）

用于計算梯度的激活值（Activations）

訓練過程中的內存動態(tài)變化

使用PyTorch的分析工具，我們可以清楚地看到，內存的使用并不是一成不變的，而是在訓練過程中不斷變化。例如，在一個訓練步驟中，內存的使用會隨著前向傳播、反向傳播和優(yōu)化步驟而起伏。

前向傳播：當我們開始訓練時，模型的激活值會迅速增加，因為模型需要計算輸入數據的輸出結果。

反向傳播：隨著梯度的計算，內存中的梯度值會逐漸積累，而用于計算梯度的激活值則會被逐步清理。

• 優(yōu)化步驟：最后，在更新模型參數時，我們需要使用所有梯度，并更新優(yōu)化器的狀態(tài)。

有趣的是，第一次訓練步驟與其他步驟看起來很不一樣。這是因為PyTorch的內存分配器在第一步中做了很多準備工作，為后續(xù)步驟預留內存，從而避免在后續(xù)步驟中頻繁搜索空閑內存塊。這種機制雖然優(yōu)化了訓練效率，但也可能導致一個常見問題：第一次訓練步驟成功了，但后續(xù)步驟卻因為內存不足而失敗。這正是因為優(yōu)化器狀態(tài)在第一次步驟后開始占用更多內存。

如何估算模型的內存需求？

為了更好地管理內存，我們需要了解模型的各個部分分別需要多少內存。首先，我們來看看模型的權重、梯度和優(yōu)化器狀態(tài)。對于一個簡單的Transformer模型，其參數數量可以通過以下公式估算：其中，h是隱藏層維度，v是詞匯表大小，L是模型層數。從這個公式可以看出，當隱藏層維度增大時，h2項會迅速增長，成為內存占用的主要部分。

接下來，我們來計算這些參數的內存需求。在傳統(tǒng)的全精度（FP32）訓練中，每個參數和梯度需要4字節(jié)，而優(yōu)化器（如Adam）還需要額外存儲動量和方差，每個參數再占用8字節(jié)。

如果使用混合精度訓練（如BF16），雖然計算速度更快，但內存需求會略有變化。在混合精度訓練中，我們通常使用BF16（2字節(jié)）進行計算，同時保留一份FP32（4字節(jié)）的模型權重副本，以確保數值穩(wěn)定性。因此，總內存需求如下：

• 模型權重（BF16）：
• 梯度（BF16）：
• 模型權重副本（FP32）：
• 優(yōu)化器狀態(tài)（FP32）：

激活值內存：訓練中的“內存大戶”

在訓練大型語言模型時，激活值內存（Activation Memory）是一個非常關鍵且復雜的部分。它不僅取決于模型的結構，還與輸入數據的長度和批次大小密切相關。激活值內存的管理，往往決定了我們能否在有限的硬件資源上訓練更大規(guī)模的模型。

激活值內存的大小并不固定，它會隨著輸入序列的長度和批次大小的變化而變化。經過仔細分析Transformer模型的反向傳播過程，我們可以估算出激活值內存的大小。具體公式如下：

其中：L 是模型的層數；seq 是輸入序列的長度；bs 是批次大??；h 是模型的隱藏層維度；n_heads 是多頭注意力機制中的頭數。

激活值重計算：用計算換內存的“魔法”

在傳統(tǒng)的訓練過程中，我們會存儲每一步的隱藏狀態(tài)（即激活值），以便在反向傳播時用來計算梯度。但有了激活值重計算，我們只需要在模型的關鍵節(jié)點存儲少量激活值，丟棄其余的激活值，并在反向傳播時從最近的保存點重新計算它們。這就好比在旅行中，你只在關鍵地點留下標記，而不是記錄整個路線的每一個細節(jié)。激活值重計算有幾種不同的策略，每種策略在內存節(jié)省和計算成本之間都有不同的權衡：1. 完整重計算（Full Recomputation）這種策略會在Transformer模型的每一層之間存儲激活值。因為它需要在反向傳播時重新執(zhí)行每一層的前向傳播，所以計算成本最高，但節(jié)省的內存也最多。通常，這種策略會使計算成本和時間增加30-40%，效果非常明顯。2. 選擇性重計算（Selective Recomputation）

選擇性重計算是一種更高效的策略。研究人員發(fā)現(xiàn)，注意力機制的計算通常占用大量內存，但重新計算的成本較低。因此，我們可以丟棄這些激活值，而只存儲前饋網絡（Feedforward）的激活值。例如，在一個1750億參數的GPT-3模型中，這種策略可以減少70%的激活值內存占用，而計算成本僅增加2.7%。

梯度累積：用“微批次”突破內存限制

梯度累積的核心思想非常簡單：將一個大批次拆分成多個小批次（微批次），然后依次對每個微批次執(zhí)行前向傳播和反向傳播，計算梯度，并將這些梯度累加起來。最后，我們用累加后的梯度平均值來更新模型參數。這樣，我們就可以在不增加內存占用的情況下，有效地模擬大批次訓練的效果。

具體來說，我們把每個前向傳播的批次大小稱為“微批次大小”（Micro Batch Size, mbs），而把兩次優(yōu)化步驟之間的總批次大小稱為“全局批次大小”（Global Batch Size, gbs）。如果我們每執(zhí)行8次前向/反向傳播后進行一次優(yōu)化步驟，那么全局批次大小就是微批次大小的8倍。雖然梯度累積解決了內存問題，但它也有一個明顯的缺點：每次優(yōu)化步驟需要多次連續(xù)的前向/反向傳播，這會增加計算開銷，從而減慢訓練速度。不過，如果你仔細思考，你會發(fā)現(xiàn)這些前向/反向傳播是可以并行化的——每個微批次的計算是獨立的，唯一的區(qū)別是輸入樣本不同。這意味著，我們可以通過多GPU并行計算來加速這一過程。數據并行：用多GPU加速模型訓練在數據并行中，每個GPU處理一個獨立的微批次數據，因此每個GPU上計算出的梯度是不同的。為了保持所有模型實例的一致性，我們需要在反向傳播過程中，通過一個稱為“全歸約”（all-reduce）的操作來平均這些梯度。全歸約是數據并行中的第一個“分布式通信原語”，它負責在GPU實例和節(jié)點之間同步和通信。

優(yōu)化數據并行的三種策略

1. 優(yōu)化一：計算與通信重疊在簡單的數據并行實現(xiàn)中，我們通常需要等待反向傳播完成后，才開始同步梯度。但其實，我們可以將通信與計算重疊起來，讓它們同時進行。具體來說，當反向傳播到達最后一層時，我們可以立即開始同步這些層的梯度，而不需要等待前面層的梯度計算完成。這樣，大部分全歸約操作可以在反向傳播過程中完成，從而提高效率。在PyTorch中，我們可以通過為每個參數添加一個全歸約鉤子函數來實現(xiàn)這一點。當某個參數的梯度計算完成后，立即觸發(fā)全歸約操作，而其他參數的梯度計算仍在繼續(xù)。這樣可以顯著減少等待梯度同步的時間。2. 優(yōu)化二：梯度分桶（Bucketing）GPU在處理大型張量時通常比處理多個小型張量更高效，通信操作也是如此。因此，我們可以將梯度分組到“桶”中，然后對每個桶中的梯度執(zhí)行一次全歸約操作，而不是對每個梯度單獨執(zhí)行全歸約。這就像打包物品時，發(fā)送幾個大箱子比發(fā)送許多小箱子更高效。通過這種方式，我們可以顯著減少通信開銷，加快通信速度。3. 優(yōu)化三：與梯度累積結合

我們之前提到，梯度累積通過多次前向和反向傳播來累積梯度，然后在最后一步更新參數。當我們將梯度累積與數據并行結合時，需要注意梯度同步的時機。在簡單的實現(xiàn)中，每次反向傳播后都會觸發(fā)全歸約操作，但這其實是不必要的。我們可以在最后一步統(tǒng)一觸發(fā)全歸約操作，從而減少通信開銷。

數據并行與梯度累積的結合

在訓練大型語言模型時，全局批次大?。℅BS）是一個關鍵參數，它直接影響模型的收斂速度和訓練效率。通過引入數據并行（Data Parallelism, DP）和梯度累積（Gradient Accumulation, GA），我們可以更靈活地調整全局批次大小，同時優(yōu)化訓練速度和硬件利用率。全局批次大小 (GBS) = 微批次大小 (MBS) × 梯度累積步數 (GA) × 數據并行實例數 (DP)在實際應用中，人們通常會優(yōu)先選擇最大化數據并行的節(jié)點數量（DP），因為數據并行是并行化的，而梯度累積是順序的。只有在數據并行無法滿足目標全局批次大小時，才會增加梯度累積的步數。例如，當我們有足夠的GPU時，可以通過增加數據并行的節(jié)點數量來加速訓練，而不是單純依賴梯度累積。

確定目標全局批次大?。℅BS）：通過查閱文獻或實驗測量模型的收斂情況，確定最佳的全局批次大?。ㄒ詔oken為單位）。

選擇序列長度：根據文獻或實驗選擇合適的訓練序列長度。通常，2-8k tokens是一個可靠的選擇。

確定單個GPU的最大微批次大?。∕BS）：逐步增加微批次大小，直到單個GPU的內存不足。

1. 確定可用的GPU數量：根據目標數據并行實例數（DP），計算所需的梯度累積步數。全局批次大小除以數據并行實例數，即為剩余的梯度累積步數。

ZeRO：零冗余優(yōu)化器，讓內存管理更高效

DeepSpeed的ZeRO（Zero Redundancy Optimizer）通過將這些組件分散到不同的數據并行節(jié)點上，顯著減少了內存冗余，同時仍然允許使用完整的參數集進行計算。接下來，我們將深入了解ZeRO的三個階段：ZeRO-1、ZeRO-2和ZeRO-3。

ZeRO-1：優(yōu)化器狀態(tài)分區(qū)

在傳統(tǒng)的數據并行中，所有節(jié)點在反向傳播后收集相同的梯度，并同時執(zhí)行相同的優(yōu)化器步驟。這不僅效率低下，還浪費了大量內存。ZeRO-1通過將優(yōu)化器狀態(tài)分成N_d（數據并行度）等份來解決這個問題。每個模型副本只保留1/N_d的優(yōu)化器狀態(tài)，并在優(yōu)化步驟中只更新1/N_d的浮點參數。

ZeRO-2：添加梯度分區(qū)

在ZeRO-1的基礎上，ZeRO-2進一步將梯度也進行分區(qū)。由于每個副本只需要與優(yōu)化器狀態(tài)對應的梯度片段，因此在反向傳播時，我們只需要執(zhí)行“歸約分散”（reduce-scatter）操作，而不是“全歸約”（all-reduce）。這樣，每個副本只需要保留1/N_d的梯度，從而節(jié)省更多內存。

ZeRO-3：添加參數分區(qū)

ZeRO-3（也稱為FSDP，即“完全分片數據并行”）是ZeRO的最終階段，它將模型參數也進行了分區(qū)。這意味著每個副本在需要時才會動態(tài)收集所需的參數片段，并在使用后立即釋放它們。這種“按需收集”的方式進一步減少了內存占用。

張量并行：打破內存限制的“魔法”在訓練大型語言模型時，激活值內存往往會成為瓶頸，尤其是在單個GPU無法容納模型的單層時。這時，張量并行（Tensor Parallelism, TP）就派上了用場。張量并行不僅分區(qū)了模型的參數、梯度和優(yōu)化器狀態(tài)，還分區(qū)了激活值，且無需在計算前收集所有片段。

張量并行的基本原理

張量并行的核心在于矩陣乘法的數學性質。在神經網絡中，矩陣乘法通常表示為 X × W，其中：

• X是輸入或激活值；
• W是神經網絡層的權重。

張量并行利用了矩陣乘法的兩個基本性質：

按列分區(qū)：可以將矩陣 B 的每一列分別與矩陣 A 相乘，然后組合結果。

按行分區(qū)：可以將矩陣 A 的每一行分別與矩陣 B 相乘，然后將結果相加。

按行分區(qū)（Row-wise Sharding）

按行分區(qū)是另一種實現(xiàn)方式。具體步驟如下：

1. 按行分割權重矩陣：將權重矩陣 W 按行分割成多個片段，并將這些片段分配到不同的GPU上。
2. 分割輸入矩陣：將輸入矩陣 X 分割成多個片段（需要一個“分散”操作，scatter）。
3. 計算局部結果：每個GPU計算局部輸入矩陣與局部權重矩陣的乘積。

匯總結果：通過一個“全歸約”（all-reduce）操作，將所有GPU上的局部結果相加，得到最終結果。

多頭注意力機制（MHA）中的張量并行

多頭注意力機制是Transformer模型的核心部分，它涉及多個矩陣乘法操作（Q、K、V）。張量并行在MHA中的應用也非常直觀：

按列分區(qū)（Column Parallel）：將Q、K、V矩陣按列分割，并分配到不同的GPU上。每個GPU計算一個或多個注意力頭的輸出。這種方法非常適合多頭注意力機制，因為每個GPU可以獨立計算一個或多個頭的注意力結果。

1. 按行分區(qū)（Row Parallel）：對于輸出投影（Output Projection），可以按行分割權重矩陣，從而減少每個GPU上的內存需求。

序列并行的基本原理

序列并行的核心思想是：在張量并行的基礎上，進一步將激活值和計算分割到不同的GPU上，但這次是沿著輸入序列的維度，而不是隱藏維度。這種方法特別適用于那些需要完整隱藏維度的操作，如LayerNorm和Dropout。

以LayerNorm為例，它需要完整的隱藏維度來計算均值和方差：

其中，均值和方差是在隱藏維度上計算的。盡管LayerNorm的計算成本較低，但它仍然需要大量的激活值內存，因為需要完整的隱藏維度。序列井行通過在序列維度上分割激活值，將內存負擔分散到多個GPU上，從而顯著減少了每個GPU的內存需求。

Diving in the GPUs – fusing, threading, mixing

Fused Kernels

Flash Attention 1-3

Mixed Precision Training