一、結論寫在前面

論文來自Microsoft。

微調大語言模型是用戶嘗試將其適應特定應用的流行選擇。然而，微調這些模型是一項艱巨的任務，因為用戶必須考慮多種因素，如資源預算、運行時間、模型大小和上下文長度等。一個具體挑戰是微調過程對內存的需求很大，這限制了所需的硬件內存和能夠處理的訓練數據上下文長度。

論文分享了針對不同微調場景的各種微調優化的詳細研究。特別是，論文評估了梯度檢查點（Gradient Checkpointing）、低秩適應（Low Rank Adaptation）、DeepSpeed的ZeRO冗余優化器（ZeRO Redundancy Optimizer）和Flash Attention。

• 論文專注于內存和運行時間，研究了不同優化組合在微調階段對GPU內存使用和執行運行時的影響。
• 論文提供了針對不同模型大小平衡內存和運行時的最佳默認優化建議。論文分享了有效策略，用于微調擁有數十億或數百億參數的非常大型模型，并在微調過程中實現大型上下文長度。
• 此外，論文提出了在GPU資源限制下進行微調的適當優化組合。

論文的結果表明：

• ZeRO-2 + LoRA是一個可靠的優化默認選項，通常能提供內存使用和微調運行時之間的最佳平衡。
• 通過使用正確的優化組合，如ZeRO-3 + LoRA + GC，論文成功微調了高達Falcon-180B的模型。
• 鑒于在GPU受限環境下微調LLM的需求，論文引入了跨模型大小x上下文長度的優化矩陣。該優化矩陣可以指導用戶選擇適合特定用例的正確優化組合。
• 論文得出結論，像DeepSpeed ZeRO這樣的優化對于微調數十億參數模型至關重要，因為它們支持數據和模型并行，以及CPU卸載。
• 結合DeepSpeed ZeRO，選擇額外的優化以避免內存不足失敗同時平衡微調運行時變得必要。

二、論文的簡單介紹

2.1 論文的背景

遷移學習為專化大型語言模型（LLMs）到特定任務或領域提供了一種非常有效的方法。在微調過程中，預訓練的LLMs通常在任務特定的數據集上進一步訓練，使其知識適應特定任務。盡管這是一種強大的技術，但由于其巨大的內存需求，微調具有數十億可訓練參數的LLMs仍然具有挑戰性。例如，在全浮點精度（32位）下微調一個10億參數模型，每個GPU大約需要24GB的高帶寬內存（HBM）。因此，在沒有任何內存優化的情況下，LLM微調過程中經常會遇到內存不足的失敗。

近年來，為優化GPU內存使用，提出了若干關鍵技術。然而，這些技術對許多用戶來說理解復雜，且其應用往往需要通過實驗來找到避免內存溢出錯誤的正確組合。這不僅需要人力資源（用于研究和實施這些優化），還需要計算資源（大量的GPU小時）來為特定任務選擇合適的優化策略。

為了簡化用戶的微調過程，諸如Azure Machine Learning等平臺提供了先進的GPU內存優化技術，使得即使使用少量GPU也能進行大型語言模型（LLM）的微調。

論文詳細介紹了流行的LLM微調優化技術，深入探討了內存與運行時間的權衡，提供了選擇最佳優化配置的指導，并分享了論文的實驗結果，這些結果可用于校準優化默認設置。

2.2 優化技術概述

論文簡要概述了四種優化技術，這些技術可用于減少微調過程中的內存瓶頸。所有這些優化技術都可以被用戶在其微調任務中，如Azure Machine學習和Hugging Face等平臺上使用。

2.2.1 Gradient Checkpointing

梯度檢查點（GC）（Chen et al., 2016）通過不在深度神經網絡（DNN）前向傳播（FP）過程中保留所有計算出的激活來審慎地利用GPU內存。相反，它在反向傳播過程中重新計算許多激活，這有助于節省GPU內存。最節省內存的策略是在每n步保存檢查點，其中n是DNN的層數（深度）。這種策略確保計算時間仍與DNN的深度成線性關系，同時將GPU內存需求從線性減少到DNN深度的平方根。因此，梯度檢查點使得能夠微調比僅增加20%-30%微調時間更大的LLMs。

2.2.2 LoRA

低秩適應（LoRA）通過減少可訓練參數的數量，從而降低微調過程中對GPU內存的需求。它通過凍結預訓練模型的權重，并在Transformer架構的每個選定層中注入可訓練的秩分解矩陣來實現這一目標。可訓練參數的數量減少了幾個數量級，降低了微調的成本，同時保持了結果的質量。例如，使用LoRA進行微調，將秩設置為64，一個700億參數模型的可訓練參數數量減少到約1.31億參數（-0.19%的原始模型大小）。

2.2.3 DeepSpeed

DeepSpeed的ZeRO冗余優化器（ZeRO）是一種內存優化技術，它提供了模型和數據并行性的好處，同時緩解了兩者的局限性。ZeRO驅動的數據并行性（Zero-DP）將模型狀態——參數、梯度和優化器狀態——跨數據并行進程進行分區，并使用動態通信調度在進程之間共享必要的模型狀態。

ZeRO-DP提供了三個優化階段，這些階段相繼提供越來越大的內存減少，同時產生一些運行時開銷。基于ZeRO-DP的微調可以分別實現第1階段和第2階段的4到8倍的內存減少，而對于第3階段（假設有大量GPU）則可以達到線性內存減少。這確實是以運行時為代價的，尤其是對于ZeRO-DP第3階段。

ZeRO優化通過包含ZeRO-Offload和ZeRO-Infinity（得到進一步增強，它們分別將優化器狀態和模型參數卸載到CPU。雖然ZeRO-Offload在所有三個ZeRO-DP階段都可用，但ZeRO-infinity僅適用于第3階段。如果不使用DeepSpeed ZeRO這樣的優化，許多大型模型（數十億參數）實際上不可能僅使用少數幾個V100或A100 GPU進行訓練。    

2.2.4 快速注意力（Flash Attention）

快速注意力幫助實現注意力計算的復雜度從與序列/上下文長度相關的二次復雜度降為線性復雜度。它利用分塊和重計算技術顯著加速注意力計算。它明智地使用靜態隨機存取存儲器（SRAM），這是GPU中最昂貴但容量最小的存儲單元，最小化SRAM與高帶寬存儲器（HBM）之間的讀寫操作。在快速注意力版本2中，進一步優化了減少較慢的非矩陣乘法操作，并沿著序列長度維度（除了批次和頭數維度外）并行化前向和后向傳播。

這些優化中的一個關鍵特點是它們彼此正交。因此，它們可以組合起來為用戶帶來漸增的內存和運行時效益。

2.3 理論分析GPU內存需求

本節的目標是幫助讀者對影響訓練過程中GPU內存消耗的計算有一個基本的理解。以下三個變量主要影響訓練期間的GPU內存：

A. 模型狀態 - 包括模型參數、梯度和優化器狀態

B. 激活 - 表示中間計算結果

C. 臨時緩沖區和碎片

模型狀態內存：論文使用Rajbhandari等人（2020）中的模型狀態內存計算來估計(A)。這里論文簡要討論圖1中所示的模型狀態內存計算。在混合精度訓練中使用Adam優化器時，可訓練參數、動量和方差以全精度（32位）保持。假設有0個參數，內存總計為(0 + 0 + 0) x 4 = 120字節。因此，K-12是Adam優化器特有的常數優化器狀態：可訓練參數 + 動量 + 方差梯度模型參數    

圖1：使用Adam優化器在混合精度設置下微調0參數模型時的模型狀態內存。模型狀態包括優化器狀態、梯度和模型參數。在未啟用任何優化的情況下，總模型狀態內存總計達到160字節

在內存計算中使用的乘數項。此外，在前向和后向傳播過程中，參數和梯度以半精度維護，這導致額外的內存需求為20 + 20 = 40字節。ZeRO-DP階段決定了哪些模型狀態（參數、梯度、優化器狀態）在GPU之間分區，從而節省內存。

激活內存：在沒有模型并行性的情況下，每個transformer層的激活內存可以使用論文第4節中的公式1（Korthikanti等人，2022）進行估計。論文依賴相同的計算來確定（B）激活內存需求的近似上限。每層的激活內存可以表示為：

其中，s是序列長度，b是微批次大小，h是隱藏維度，a是注意力頭的數量。總激活內存計算為transformer層數 × 每層激活內存。

表1：ZeRO-DP三個階段理論與實測GPU內存分配（以GB為單位）的比較。微調模型為LLaMA-2-7B，使用8個A100 GPU（每個具有80 GB HBM）。理論估計可用于在運行微調作業之前近似實測內存需求    

表 

 展示了論文在比較理論上預期的和經驗上分配給LLaMA-2-7B微調的GPU內存方面的結果，涉及ZeRO-DP的不

同階段。論文使用ZeRO-1、ZeRO-2和ZeRO-3分別表示ZeRO-DP的第1、2和3階段。理論內存是通過匯總模型狀態和激活的內存來計算的。論文使用以下模型默認參數：transformer層數=32，隱藏維度大小h=4096，注意力頭數a-32，理論上計算出總激活內存為1.48 GB。經驗GPU內存指的是微調過程中分配的峰值GPU內存。論文展示了有無CPU卸載優化器狀態的經驗內存。

論文的結果表明，理論計算可以作為微調過程中GPU內存消耗（無CPU卸載）的粗略估計。在運行微調任務之前計算這一估計有助于預測內存需求并有效規劃資源分配。論文注意到，啟用CPU卸載優化器狀態可以將內存使用量減少至不啟用卸載時的四分之一。因此，CPU卸載可以允許對具有數十億參數的模型進行微調時，對GPU和CPU資源進行高效利用，這是一個內存密集型過程。

需要注意的是，在沒有ZeRO-DP優化的情況下，使用傳統數據并行性對70億參數模型進行完全微調將需要超過112（16 x 7）GB的GPU內存。這樣的內存需求遠遠超過了A100s提供的80 GB GPU HBM。

2.4 實驗

論文首先分析微調優化對內存和運行時的影響。論文的分析結果用于推薦一組優化，這些優化可以作為平衡默認值，以優化微調期間的內存和運行時。隨后，論文深入研究大型模型（具有數十億參數）的微調，并探索使這種微調成為可能的優化。接著，論文檢查Flash Attention 2在LLMs微調長上下文數據中的作用。最后，論文研究在資源受限的GPU（特別是V100s）上進行微調的具體情況。論文的目標是識別能夠在各種模型大小和上下文長度下，甚至在有限資源下實現高效微調的優化。

2.4.1 設置

論文在Causal Language Modeling任務上對LLaMA-2（7B、13B、70B）和Falcon（180B）系列的模型進行微調。以下是論文的實驗設置的詳細信息：

• 數據：所有實驗中使用了Sampled Samsum數據集（Gli et al., 2019）進行微調。
• 計算：使用了標準ND4Ors v2（8xV100）（Azure, 2022）和標準ND96amsr A100 v4（8xA100）（Azure, 2024）作為GPU計算資源。
• 優化器：所有實驗使用混合精度設置和AdamW優化器（β1=0.9，β2=0.99），配合線性調度器和學習率4e-4。
• 序列長度、批次大小和周期：序列長度通常設置為256。在探索序列長度影響的研究中，通過填充擴展輸入序列的長度。所有實驗中使用的有效批次大小為8。注意，可以通過梯度累積來增加有效批次大小，同時實現對可用GPU內存的高效利用。所有模型都進行單周期微調。
• LoRA：在使用LoRA的實驗中，LoRA的rank設置為64，alpha設置為32。
• CPU卸載：除非明確指出，論文使用ZeRO-Offload在微調過程中將優化器狀態和計算卸載到CPU。這有助于使用可用的GPU和CPU的計算及內存資源訓練數十億參數的模型。

在論文的實驗中，GPU內存使用量通過微調過程中所有使用的GPU上分配的峰值GPU內存來衡量。

2.4.2 最佳默認優化以平衡內存和運行時間

論文已經確定，對于相對較小的7B參數模型，即使使用A100，不進行任何優化進行微調也是不可行的。因此，論文使用DeepSpeed ZeRO，它提供了模型和數據并行性。論文檢查了在微調框架中應默認啟用的微調優化集合。論文選擇這些優化的標準是基于它們在ZeRO-DP結合下，能夠在內存使用和運行時間之間達到最佳平衡的能力。

論文在一個節點（8x A100 GPU）上對LLaMA-2-7B進行了五種不同配置的微調：(a) 未啟用優化 (b) 啟用梯度檢查點 (c) 啟用LoRA (d) 啟用Flash Attention 2 (e) 未啟用CPU卸載。每種配置都啟用了所有三個ZeRO-DP階段。論文選擇A100 GPU進行這些實驗，因為它們具有80 GB的大內存容量。這種充足的內存使論文能夠通過開啟和關閉各種優化來靈活地進行實驗，同時避免內存不足的失敗。    

圖2：LLaMA-2 7B在不同優化配置下，使用ZeRO-1、ZeRO-2和ZeRO-3時的GPU內存使用情況和微調運行時間。使用ZeRO-2結合LoRA提供了最佳的內存使用與運行時間平衡

圖2展示了每種配置和ZeRO-DP階段下的GPU內存使用和運行時間。根據結果，論文得出以下結論：

1.從運行時間和內存使用的角度來看，最佳的優化組合是配置(c)，即ZeRO-DP + LoRA。由于LoRA減少了總的可訓練參數，它節省了內存并縮短了運行時間。在Zero-DP + LoRA配置中，ZeRO-1和ZeRO-2階段的內存使用及運行時間大致相似。當使用LoRA（rank=64）時，可訓練參數減少到約3360萬個，對于7B模型而言。優化器狀態和可訓練參數的梯度占據了GPU內存的極小部分。在ZeRO-1和ZeRO-2中，內存消耗主要由總模型參數（70億）主導，這些參數在半精度下占用14GB內存。

1.配置(a)和(d)在ZeRO-DP階段的表現符合預期，即從階段1到2到3，內存消耗下降，運行時間增加。對于配置(b)和(e)，ZeRO-1的運行時間略高于ZeRO-2，內存使用遵循隨著Zero-DP階段增加而減少的預期趨勢。

2.在(e)中禁用CPU卸載提供了最快的運行時間，但需要2-4倍更多的GPU內存。

論文得出結論，ZeRO-2與LoRA的結合是維持內存使用和運行時間平衡的優秀默認選擇。此外，ZeRO-2能處理比ZeRO-1更大的模型，這有助于支持更廣泛的模型尺寸范圍。對于某些用例，如復雜的問答任務或構建聊天機器人，用戶可能更喜歡完全微調模型。因此，是否使用LoRA由用戶自行決定。因此，論文建議在微調過程中默認激活ZeRO-2。    

2.4.3 大型模型的微調

微調大型模型如LLaMA-2 70B或Falcon 180B時，如果沒有啟用正確的優化設置，常常會遇到內存不足（OOM）錯誤。在這種情況下，啟用ZeRO-3是一個可行的替代方案，盡管它會增加運行時間。論文來檢查在ZeRO-3中，經過個節點和每個節點個GPU微調后的模型狀態內存項：

這里，表示總可訓練參數（以十億計）。假設一個標準的配置，每個節點有

個GPU，這將隨著模型尺寸的增加而增加。例如，在個節點上完全微調一個

億參數的模型，將占用40GB的模型狀態內存，這可以在具有80GB HBM的A100上容納。在實踐中，具有數百億參數的大型模型通常與LoRA一起微調，并使用ZeRO-3與CPU卸載。這進一步降低了內存需求，使其達到可管理的規模。             

論文能夠在單個節點上使用8xV100 GPU（32 GB HBM）結合ZeRO-3 + LoRA對LLaMA-2 70B進行微調。表顯示，在論文的實驗中，70B模型的實際GPU內存消耗約為15.54 GB。此外，可以看出，對于LLaMA-2 7B和13B這樣的較小模型，使用ZeRO-3 + LoRA會導致GPU內存未充分利用。因此，論文的研究表明，對于微調具有數十億或數百億參數的大型模型，啟用ZeRO-3是必不可少的。此外，當與LoRA結合使用時，實際的總體內存需求顯著減少。

表2：在8xV100 GPU上使用ZeRO-3 + LoRA對LLaMA-2模型進行微調的GPU內存使用量和時間

對于較小模型，ZeRO-3 + LoRA顯然過于冗余，因為它可能導致GPU未充分利用以及運行時間增加。    

2.4.4 長上下文微調

              圖3：LLaMA-2 70B在不同上下文長度下，使用和不使用Flash-Attention 2對GPU內存使用和微調時間的影響。在A100上啟用Flash-Attention 2顯著降低了較大上下文長度（如4096)的內存消耗和運行時間

上下文長度是微調LLMs時的關鍵因素。特別是在使用包含非常長文本序列的數據集進行微調時尤為如此。沒有Flash-Attention 2（FA2），注意力計算隨著序列長度的平方增長而增加。由于V100 GPU不支持FA2，因此在使用較長上下文長度時，對GPU內存消耗的影響可能很大。為了研究在微調過程中長上下文長度的影響，論文繪制了在A100（啟用和未啟用FA2）和V100（未啟用FA2）上不同上下文長度的內存和運行時間。

具體來說，論文在8xV100和8xA100 GPU上分別使用1024、2048和4096的上下文長度對LLaMA-2 70B模型進行微調。在微調過程中啟用了LoRA和梯度檢查點技術，并基于ZeRO-3進行優化。

圖3展示了Flash-Attention 2對GPU內存和運行時間的影響。論文的結果顯示，隨著上下文長度增加，至4096時，使用Flash-Attention 2相比于不使用，GPU內存使用和運行時間有顯著減少。如圖3所示，在兩塊A100 GPU上或分別在A100和V100 GPU上比較使用和不使用Flash-Attention 2的運行時，這一趨勢是一致的。

論文的結論如下：高性能GPU如A100，由于其更大的HBM，可以在沒有Flash-Attention 2的情況下支持相對較長的上下文微調。然而，在兼容的GPU架構上激活Flash-Attention 2是最優的，因為它允許額外的內存節省和減少運行時間。當使用允許較大上下文長度微調的模型訓練長文本數據時，這一點變得尤為重要。    

2.4.5 在GPU資源限制下的微調

圖4：使用V100 GPU時，不同大小LLMs微調的最佳配置。所有LLaMA-2實驗均使用8x V100進行，而Falcon-180B需要16xV100。由于V100不支持Flash-Attention 2，因此未包含在內

本節討論在資源有限的情況下對LLMs（參數范圍從7B到180B）進行微調。論文將資源限制定義如下：

?可用的低資源GPU具有有限的HBM。

?可用的GPU數量較少。

論文對LLaMA-2 7B、13B、70B和Falcon-180B模型在五種不同的上下文長度上進行了微調。遵循資源約束規范，論文的實驗在一臺標準ND4Ors v2（8xV100 GPU）上進行，配備32 GB HBM。唯一例外的是Falcon-180B模型，由于其龐大的規模，論文在微調過程中使用了兩個節點（16x V100 GPU）。圖4展示了為滿足GPU內存需求同時優化運行時間應啟用的首選優化集合。

需要注意的是，雖然可能存在其他允許成功微調的組合，但論文選擇了那些最小化微調時間的組合。例如，用戶可以選擇啟用所有優化，但通常并非所有優化都是必要的，且可能會不必要地增加微調時間。

論文的結果中的關鍵見解包括：

• 即使是為了在有限數量的GPU上適配像7B參數這樣的小型模型，結合模型和數據并行是先決條件，除非論文利用量化技術。
• 隨著上下文長度或模型大小的增加，轉向更高階段的ZeRO-DP（零冗余優化器-數據并行）以促進微調變得重要。這一趨勢在表格中清晰可見，從左到右或從上到下的移動對應于更高階段的ZeRO-DP。
• 梯度檢查點（GC）是一種有效的內存節省優化技術，尤其適用于大型模型。如圖4所示，為LLaMA-2 13B和70B啟用GC使得支持的上下文長度從512擴展到4096，否則這是不可能的。類似地，GC對于在論文檢查的所有上下文長度上啟用Falcon-180B的微調至關重要。
• 盡管Flash-Attention 2在V100 GPU上不支持，但論文的實驗結果表明，在支持的架構上應始終啟用它。值得一提的是，圖中概述的配置可用于在具有比V100更高HBM的GPU上進行微調。雖然這些配置可能不是所有GPU的最佳選擇，但論文的發現可以通過進一步的研究擴展到其他GPU。

論文標題：A Study of Optimizations for Fine-tuning Large Language Models

論文鏈接：???https://arxiv.org/pdf/2406.02290??    

本文轉載自 ??AI帝國??，作者： 無影寺