自拍偷在线精品自拍偷,亚洲欧美中文日韩v在线观看不卡

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

發(fā)布于 2024-7-4 08:48
瀏覽
0收藏

一、背景

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

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

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

二、引言

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

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

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

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

在 Bloom [2211.05100] BLOOM: A 176B-Parameter Open-Access Multilingual Language Model 的預訓練中作者也頻繁遇到各種異常，因此會每 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)計了一個代表性任務在一個月內(nèi)遇到的錯誤。數(shù)據(jù)顯示，由于這些錯誤，該作業(yè)在一個月內(nèi)失敗了 40 次：

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

2.2 Checkpoint 大小

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

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

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

以此可以推測出 LLM 預訓練 Checkpoint 的大?。海℅PT 175B 模型，Adam 優(yōu)化器）：

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

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

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

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

三、Checkpoint 壓縮（ExCP）

3.1 方法

LLM 預訓練時間往往比較長，可能會持續(xù)幾個月。為了應對故障等問題會頻繁的保存 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ù)量的模型上進行評估，410M 模型可以壓縮 70x 大小，而基本不損失下游任務上的精度。

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

3.2 結(jié)果

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

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

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

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

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

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

四、Checkpoint 快速保存（FastPersist）

4.1 方法

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

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

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

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

4.2 結(jié)果

作者進行了一系列的實驗來評估 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，最高達到了 146GB/s，將近是 8 個節(jié)點理論峰值的 80%。
  • Figure 9c：使用 FastPersist，模型的訓練速度也明顯提升，并且模型越小提升越明顯。從 1.6x（gpt 13B）到 21.8x（gpt 0.7B）。
  • Figure 9d：進一步繪制了訓練速度與 DP 數(shù)目的關(guān)系。?

?

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

五、參考鏈接

  1. ??https://arxiv.org/abs/2205.01068??
  2. ??https://arxiv.org/abs/2204.02311??
  3. ??https://arxiv.org/abs/2311.16867??
  4. ??https://arxiv.org/abs/2211.05100??
  5. ??https://arxiv.org/abs/2403.07648??
  6. ??https://arxiv.org/abs/2401.02954??
  7. ??https://arxiv.org/abs/2406.04594??
  8. ??https://arxiv.org/abs/2406.11257??
  9. ??https://arxiv.org/abs/2406.13768??

本文轉(zhuǎn)載自??AI閑談??，作者： AI閑談 ????


贊
收藏
回復
舉報
回復
相關(guān)推薦