Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮

amei2000go

發(fā)布于 2024-7-4 08:48

瀏覽

0收藏

一、背景

在 LLM 的預(yù)訓(xùn)練過程中，通常會不斷地保存 Checkpoint，主要有以下幾個好處：

中斷恢復(fù)：在大規(guī)模訓(xùn)練中，硬件故障和偶然的軟件故障是不可避免的。通過保存 Checkpoint，可以在發(fā)生異常時恢復(fù)訓(xùn)練，避免從頭開始。
進(jìn)度管理：LLM 需要評估的指標(biāo)很多，但通常不會在模型訓(xùn)練中實(shí)時評估，而是通過保存 Checkpoint 進(jìn)行離線評估。這有助于評估模型是否朝著正確的方向發(fā)展。
增量學(xué)習(xí)：保持 Checkpoint 也為后續(xù)的增量學(xué)習(xí)和遷移學(xué)習(xí)提供基礎(chǔ)，可以從中間階段繼續(xù)訓(xùn)練。
模型調(diào)優(yōu)：在需要重新調(diào)整超參數(shù)或者跳過部分?jǐn)?shù)據(jù)（比如出現(xiàn) spike）時，可以回到之前的 Checkpoint 進(jìn)行調(diào)整。

雖然保存 Checkpoint 有眾多好處，但也會帶來一些挑戰(zhàn)。例如，頻繁保存 Checkpoint 可能影響訓(xùn)練速度，也可能帶來較大的存儲壓力。本文中簡單介紹最近兩個與 Checkpoint 相關(guān)的工作：一個是 Checkpoint 壓縮，降低存儲壓力；另一個是加快 Checkpoint 保存速度。

二、引言

2.1 LLM 預(yù)訓(xùn)練的挑戰(zhàn)

在 OPT [2205.01068] OPT: Open Pre-trained Transformer Language Models 的 57 天預(yù)訓(xùn)練中，前期由于各種問題至少手動重啟了 35 次任務(wù)。為了避免每次需要人工重啟，后續(xù)加入了自動重啟機(jī)制，由于硬件故障又觸發(fā)了 70+ 次，平均每天都要重啟一次任務(wù)。

在 PaLM [2204.02311] PaLM: Scaling Language Modeling with Pathways 的預(yù)訓(xùn)練中，作者遇到了 20 次毛刺需要重啟訓(xùn)練。

在 Falcon [2311.16867] The Falcon Series of Open Language Models 中作者也遇到了 9 次毛刺需要重啟訓(xùn)練。

在 Bloom [2211.05100] BLOOM: A 176B-Parameter Open-Access Multilingual Language Model 的預(yù)訓(xùn)練中作者也頻繁遇到各種異常，因此會每 3 小時保存一次 Checkpoint。

在上海 AI lab 和商湯的 [2403.07648] Characterization of Large Language Model Development in the Datacenter，作者也有同樣的問題，其每 30 分鐘就保存一次 Checkpoint。

在 DeepSeek [2401.02954] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism 中，作者更是每 5 分鐘就保存一次 Checkpoint。

如下圖 Table 1 所示，阿里 C4 [2406.04594] Boosting Large-scale Parallel Training Efficiency with C4: A Communication-Driven Approach 中，作者統(tǒng)計了一個代表性任務(wù)在一個月內(nèi)遇到的錯誤。數(shù)據(jù)顯示，由于這些錯誤，該作業(yè)在一個月內(nèi)失敗了 40 次：

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

2.2 Checkpoint 大小

在 LLM 預(yù)訓(xùn)練中，為了保證模型異常后能夠繼續(xù)訓(xùn)練，需要保存模型訓(xùn)練的中間狀態(tài)，也就是 Checkpoint，Checkpoint 中除了模型參數(shù)外，還包含模型的優(yōu)化器狀態(tài)，以及一些超參（比如 batch size，learning rate 等，通常很小，大小可以忽略）。

以 LLM 經(jīng)常使用的 Adam 優(yōu)化器為例，Adam 是一種結(jié)合了動量（Momentum）和自適應(yīng)學(xué)習(xí)率（如 RMSprop）技術(shù)的優(yōu)化算法，廣泛用于深度學(xué)習(xí)模型的訓(xùn)練。Adam 優(yōu)化器的核心優(yōu)勢在于能夠根據(jù)參數(shù)的更新歷史自適應(yīng)調(diào)整每個參數(shù)的學(xué)習(xí)率，這有助于加快收斂速度并提高訓(xùn)練穩(wěn)定性。為了實(shí)現(xiàn)這一點(diǎn)，Adam 維護(hù)了 2 個額外的參數(shù)（或“矩”），也就是每個參數(shù)都對應(yīng) 2 個優(yōu)化器參數(shù)：

一階矩 v：過去梯度的移動平均，提供了平滑過的梯度方向。這有助于增加穩(wěn)定性，使模型在參數(shù)空間中的移動更加平滑。
二階矩 m：過去梯度平方的指數(shù)移動平均，代表過去梯度的變化率，允許每個參數(shù)有自適應(yīng)的學(xué)習(xí)率。參數(shù)更新的規(guī)模由其梯度的不確定性決定，有助于調(diào)整學(xué)習(xí)步長，避免在參數(shù)空間中過大或過小的步進(jìn)。

以此可以推測出 LLM 預(yù)訓(xùn)練 Checkpoint 的大小：（GPT 175B 模型，Adam 優(yōu)化器）：

模型參數(shù)：175B 參數(shù) * 2 Byte = 350 GB。在常見的混合精度訓(xùn)練中，為了減少累積誤差，提高訓(xùn)練穩(wěn)定性，通常會保留一份FP32 精度的模型參數(shù)，如果是這種情況，則對應(yīng)的模型參數(shù)為 700GB 左右。
Adam 優(yōu)化器狀態(tài)：通常使用 FP32 存儲，175B * 4 Byte * 2 = 1400 GB，對應(yīng)上述的一階矩和二階矩。

綜上，LLM 預(yù)訓(xùn)練中顯存占用大概是 10 倍參數(shù)量或 12 倍參數(shù)量。

如下圖 Table 1 所示為一系列模型對應(yīng)的 Checkpoint 大小，差不多是 13 倍左右：

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

三、Checkpoint 壓縮（ExCP）

3.1 方法

LLM 預(yù)訓(xùn)練時間往往比較長，可能會持續(xù)幾個月。為了應(yīng)對故障等問題會頻繁的保存 Checkpoint，常見的為 1-3 小時保存一次，與此同時，每個 Checkpoint 都很大，可能上百 GB 到 2TB，這會給存儲帶來極大的壓力。為了解決這個問題，常見的方式是定期、滾動地刪除更早的 Checkpoint。

在 [2406.11257] ExCP: Extreme LLM Checkpoint Compression via Weight-Momentum Joint Shrinking 中，作者提出了 Checkpoint 壓縮的方案 ExCP（Extreme Checkpoint Compression），其利用 Checkpoint 之間的冗余信息來壓縮，在 410M 到 7B 參數(shù)量的模型上進(jìn)行評估，410M 模型可以壓縮 70x 大小，而基本不損失下游任務(wù)上的精度。

ExCP 具體方案如下圖 Figure 2 所示，考慮到相近 Checkpoint 的模型權(quán)重相似度比較高，而優(yōu)化器狀態(tài)會通過滑動平均，相似度不一定特別高，因此主要是利用權(quán)重相似性來稀疏化：

左側(cè)：相鄰權(quán)重相減，求出殘差，其殘差的稀疏性比較高。
中間：

M2W：首先使用殘差權(quán)重和二階矩 m 確定裁剪 Mask Mw。

W2M：然后使用裁剪 Mask Mw 和一階矩 v 確定裁剪 Mask Mo。

右側(cè)：
使用 Mw 裁剪權(quán)重。
使用 Mo 裁剪一階矩和二階矩。
壓縮 Checkpoint。

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

PS：以上的過程考慮了兩個相鄰的 Checkpoint，那么針對整個訓(xùn)練中的 Checkpoint 如何處理呢？最簡單的想法是針對原始的 Checkpoint，每一個都依賴之前的 Checkpoint 進(jìn)行壓縮。但是這種方式會存在一個問題，從壓縮的 Checkpoint 恢復(fù)出來的只是殘差，依然需要上一個原始 Checkpoint 才能恢復(fù)當(dāng)前的 Checkpoint，然而此時原始的 Checkpoint 可能已經(jīng)被清理。因此，作者采用的是當(dāng)前 Checkpoint 與上一個已經(jīng)稀疏化處理的 Checkpoint 的殘差。這樣只用保留最原始的 Checkpoint，之后的 Checkpoint 都順序解壓即可恢復(fù)所有 Checkpoint。當(dāng)然，也可以考慮階段性地保留一些中間的原始 Checkpoint。

3.2 結(jié)果

如下圖 Table 3 所示，作者使用 Pythia-410M 進(jìn)行了驗(yàn)證，可以從原始的 4.53GB 壓縮到 0.06GB，壓縮 70+ 倍；并且在下游任務(wù)上的平均指標(biāo)只從 43.11 下降到 42.93：

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

如下圖 Figure 4 所示，在整個訓(xùn)練過程中對 Checkpoint 進(jìn)行壓縮，其 Loss 與未壓縮的原始 Loss 非常接近：

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

如下圖 Table 4 所示，作者在 PanGu-π-1B 和 PanGu-π-7B 上進(jìn)行了實(shí)驗(yàn)，同樣在壓縮 20+ 倍的情況下平均精度甚至更高了一點(diǎn)點(diǎn)：

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

四、Checkpoint 快速保存（FastPersist）

4.1 方法

在 [2406.13768] FastPersist: Accelerating Model Checkpointing in Deep Learning 中，Deepspeed 的作者提出了 FastPersist，其主要是為了加快 DL 訓(xùn)練中 Checkpoint 保存的速度，最多可以比基線加快 116x。FastPersist 主要包含 3 點(diǎn)：

加快 Checkpoint 寫入（下圖 b）：利用節(jié)點(diǎn)上的 NVMe 設(shè)備的高速寫入帶寬，并使用新的 I/O 庫（比如 Linux 中的 libaio 和 io_uring）來異步的加快寫入性能。
并行化寫入（下圖 c）：在 DP（Data Parallelism）之間分布式寫入，充分利用不同 DP 上的硬件，減少寫入延遲。
Pipeline 寫入（下圖 d）：創(chuàng)建專門的 Python 進(jìn)程來異步的寫入。

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)

PS：其實(shí)之前很多文章中也提到了加快 Checkpoint 保存，避免影響訓(xùn)練速度的工作。主要是分布式保存，異步保存等。

4.2 結(jié)果

作者進(jìn)行了一系列的實(shí)驗(yàn)來評估 FastPersist 的效果，結(jié)果如下圖 Figure 9 所示：

Figure 9a：在 128 個 V100 GPU 上，F(xiàn)astPersist 獲得了 28x（gpt 13B）到 116x（gpt 0.7B）的加速。
Figure 9b：隨著 DP（Data Parallelism）的增加，保存速度近線性增加，對于 gpt 13b，最高達(dá)到了 146GB/s，將近是 8 個節(jié)點(diǎn)理論峰值的 80%。
Figure 9c：使用 FastPersist，模型的訓(xùn)練速度也明顯提升，并且模型越小提升越明顯。從 1.6x（gpt 13B）到 21.8x（gpt 0.7B）。
Figure 9d：進(jìn)一步繪制了訓(xùn)練速度與 DP 數(shù)目的關(guān)系。?

Excp & FastPersist：數(shù)十倍 LLM Checkpoint 保存加速和壓縮-AI.x社區(qū)