大模型訓練效率太低，存儲空間太大？有一個新的開源技術能優化這個流程！

隨著大模型工程技術的迅猛進步，提升大模型訓練效率已成為推動其發展的關鍵要素。訓練效率 = 訓練吞吐 × 訓練有效率 × 收斂效率，其中，訓練有效率的保障離不開靈活且強大的模型恢復機制。據悉，Meta的萬卡集群在訓練Llama3.1時，平均每3⼩時便遭遇⼀次故障，這凸顯了完善的大模型Checkpoint保存與恢復功能對於提升訓練效率的重要性。

為應對這一挑戰，PaddleNLP大語言模型套件針對大規模訓練場景，研發並開源了大模型統一存儲技術—— Unified Checkpoint。該技術集訓推一體、異步保存、快速恢復、無損壓縮等多重優勢於一身，顯著優化了大模型訓練流程。以下是其主要功能亮點：

• 支持全分佈式策略調整自適應轉換，提升模型訓練的靈活性與可擴展性

  工業級的大模型訓練根據機器數量的變化，會涉及多種分佈式策略的靈活組合和調整。UnifiedCheckpoint通過存儲格式與分佈式策略的解耦設計，支持Checkpoint在不同分佈式策略間的自動切換，用戶無需感知這一變化，可顯著提升大模型恢復訓練的靈活性與可擴展性。

• 訓練-壓縮-推理統一存儲協議，無需手動轉換提升全流程體驗

  傳統上，訓練階段Checkpoint的參數結構與推理階段參數切分格式往往不相同，導致兩者間額外的轉換成本。PaddleNLPUnified Checkpoint實現了訓練與推理參數格式的無感切換，極大降低了兩者銜接的複雜度，進一步提升了PaddleNLP套件的便捷性。

• Checkpoint無損壓縮結合異步保存，實現秒級存儲並降低模型存儲成本

  大模型Checkpoint的存儲不僅耗時⻓，還佔⽤龐大的磁盤空間。以Llama-2-13B模型為例，單個Checkpoint的存儲耗時約150秒，磁盤空間消耗高達182GB。PaddleNLP通過異步存儲和模型Checkpoint無損壓縮算法最終可以實現秒級存儲和降低80%左右的大模型存儲空間。

一、Unified  Checkpoint統一訓推存儲，全面支持任意分佈式策略組合

大模型Checkpoint存儲技術面臨諸多挑戰，特別是對於大規模分佈式訓練而言：

• 傳統Checkpoint存儲方案⽆法適⽤複雜的大模型訓練場景。在真實的大模型訓練場景中，分佈式策略組合⾮常多且複雜，同時訓練資源的擴縮容導致Checkpoint的保存需要很好的靈活性和自適應性。例如用戶因為訓練集群縮容就需要Checkpoint恢復訓練時能適應分佈式策略帶來的變化。

• 大模型訓練和推理的Checkpoint格式不統一導致傳統Checkpoint存儲方案無法適用：與上面提到機器擴縮容類似，訓練和推理的采⽤的分佈式策略不同，導致傳統的Checkpoint方法無法將訓練保存模型參數直接拿來預測。

為瞭解決上述問題，我們提出了Unified Checkpoint統一存儲的方案。

1.1  Checkpoint存儲格式

在Checkpoint存儲格式上，我們在⻜槳框架原先的存儲協議上，針對大模型引入了Safetensors作為Unified Checkpoint的序列化保存格式。使⽤Safetensors格式具有非常多的優勢，具體可⻅下圖：

Safetensors的優勢主要有：

• 安全性（Safe）：Safetensors格式中不包含任何可以執行的代碼，確保文件加載時不會觸發任意代碼執⾏。

• 零拷⻉（Zero-copy）: Safetensors格式具備接近零拷⻉的效果，在將文件加載到內存時，通過內存映射技術將文件的數據直接映射到內存，從而可避免加載過程中多個進程重覆加載拷⻉同一文件的情況。

• 惰性加載（Lazy Loading）: 在分佈式環境下，實現張量的部分加載是⾮常方便且效率更高的。這一點也非常有利於我們實現不同分佈式策略切換的任意加載，而不需要為此創建新的權重副本，大大降低操作複雜度。

基於上述的這些特性，Unified Checkpoint的存儲格式採用了Safetensors格式。

1.2 統一Checkpoint存儲，支持不同分佈式策略切換

原始的Checkpoint存儲⽅式，在存儲時會按照實際使用的不同分佈式策略進行保存，從而導致在進行分佈式策略切換時難以靈活擴展。具體⻅下圖，固定形式的Checkpoint存儲格式使得我們難以靈活應對不同的分佈式情況，往往需要⼈工定製相應轉換腳本，導致可擴展性低。下面是原始Checkpoint格式無法擴展的示意圖。

UnifiedCheckpoint 設計了存儲與分佈式策略解耦的方案，做到Checkpoint靈活擴展的效果。其中最核心的模塊是參數路由表的更新，路由表包含了模型參數名稱、文件路徑、機器結點三個關鍵信息。在分佈式訓練中訓練節點在加載(Load)模型⽂件階段，通過路由表獲取模型參數所在節點位置以及參數文件名稱，通過網絡傳輸獲取模型參數；當分佈式策略發生發變化模型參數的分佈也隨之變化，模型參數在保存(Save)模型文件階段，各個節點需要更新參數的所在節點以及模型所在文件信息到全局路由表。

在Unified Checkpoint的設計原則中，需要確保每個參數被保存時需要是一個完整的張量，而不是被多個節點切分保存，這樣可以做到存儲下來的Checkpoint文件與分佈式策略解耦。如果同一個參數被多個節點保存，那在分佈式策略發生變化時參數的切分將會極為複雜。上圖中我們通過Gather的通信將參數進行融合，然後再將參數信息同步到路由表中。同時為了加快模型參數的保存，每個worker進程都會併發保存部分的完整張量，提升保存效率，下面是參數融合的示意圖。

在Checkpoint加載過程中，我們可以具體區分為兩種情況，分別表示原地加載和動態加載。

• 原地加載的情況，主要指當前訓練進程所需要的參數文件均在其所在機器，此時主要依賴Safetensors格式的惰性加載特點，實現零冗餘加載。每個worker 各自讀取參數文件中所需的張量部分即可，無需進行跨卡或者跨機發送。原地加載的情況已經基本可以適配大多數的訓練場景，在確保不同機器為共享存儲的情況下，則可以靈活地進行不同分佈式策略之間的切換。

• 動態加載的情況，主要指訓練進程所需要的參數文件不在當前機器，需要進行跨卡或者跨機之間的張量發送。這種情況往往發生在機器為飛共享存儲，且需要切換分佈式策略的情況。為瞭解決這一問題，我們實現了自適應動態加載機制。通過構造了張量動態路由表，用於確認哪些worker需要負責Tensor的發送，哪些進程需要負責接收Tensor。在構造完成動態路由表之後，即可進行不同卡之間的Send、Recv異步通信。在每個worker獲取到所需要的參數之後，即可完成Checkpoint的加載，實現自適應動態加載，進而實現不同分佈式策略、不同機器之間的切換。

1.3 訓練和推理模型存儲格式無縫切換

我們將模型的Checkpoint 通過 Unified Checkpoint 技術轉換為 Safetensors 格式進行保存，這種統一的存儲⽅案具備高度的靈活性和兼容性，可以無縫銜接到模型推理的不同環節，⽆論是在動態圖直接推理，還是在將動態圖轉換為靜態圖的過程中，均能實現平滑切換。這一特性消除了用戶手動處理參數合併或切分的繁瑣操作，不僅簡化了使用流程，還顯著提升了開發和部署效率，為⼤模型下遊應⽤提供了更加便捷和高效的支持。

二、Checkpoint存儲優化，存儲最快加速95%，空間最高可節省78.5%

除了統⼀訓練推理存儲格式，支持靈活切換不同的分佈式策略以外，我們還針對Checkpoint存儲速度和存儲空間進行了優化。⼀般的Checkpoint存儲⽅式往往⾯臨兩⼤挑戰：

• 存儲耗時⻓，影響訓練效率：對於擁有百億乃至千億參數的大型模型，每次Checkpoint的保存可能需要幾分鐘甚至十幾分鐘。若設置較短的保存間隔，會頻繁中斷訓練進程，不僅浪費訓練資源，還會顯著降低整體訓練效率。

• 存儲空間龐⼤，增加存儲成本：以⻜槳大模型訓練框架為例，一個原始的Checkpoint所佔用的存儲空間可能是模型權重本身所佔空間的7倍。對於70B的模型，這意味著需要接近1TB的磁盤空間來存儲一個 Checkpoint。

為瞭解決上述兩個問題，我們分別進行了針對性解決和優化，最終使得Unified Checkpoint的存儲效率相⽐原先最大可加速 95%，存儲空間最高可節省 78.5%。

2.1 異步存儲大幅減少存儲耗時95%

通常情況下，模型訓練和模型存儲是同步進行的。然而，在保存模型Checkpoint時，傳統的做法會中斷GPU訓練進程以進行模型保存，這往往導致訓練效率顯著降低。具體來說，保存Checkpoint通常包含兩個關鍵階段：

1. 顯存複製操作（D2H Copy）：將GPU顯存中的參數複製到主機內存。

2. 模型參數持久化（Save to Disk）：將內存中的模型參數通過序列化方式保存到硬盤中。

由於這兩個階段均涉及IO操作，因此存在將計算操作（訓練）與IO操作（模型保存）並行的可能性，通過重疊（Overlap）這兩種操作來減少存儲開銷。

為了實現這一目標，我們採用了異步存儲的流水線方法。通過開闢異步存儲⼦進程，我們顯著降低了主進程被阻塞的時間，從而提升了系統整體的吞吐量。值得注意的是，顯存複製操作（D2HCopy）需要在訓練進程中執行，而異步存儲子進程則需要訪問主進程中的模型參數內存。因此，在第一次進行Checkpoint保存 時，我們會分配一段連續且固定的共享內存作為主機端的存儲位置。主進程將GPU端的數據複製到這段共享內存後，異步進程會從共享內存中讀取數據，並將其持久化到磁盤。最終，PaddleNLP UnifiedCheckpoint的異步存儲架構可簡單描述如下：

在實際的存儲測試中，我們對不同規模的模型進行了存儲對⽐測試。以下表格對⽐中，保存介質均為NVMeSSD。我們目前測試情況下，發現存儲耗時最大可減少95%。

2.2 Checkpoint壓縮最高可節省 78.5 %存儲空間

在大模型訓練過程中，為了能夠在訓練中斷後從最新的Checkpoint快速恢復，我們需要不斷地保存模型訓練的Checkpoint。理論上，保存的Checkpoint頻次越⾼，恢復訓練時所需重訓的step就越少。但是大模型的 Checkpoint會佔用大量的磁盤空間，從⽽限制了能夠保存的Checkpoint數量。那麼，我們是否可以對模型的 Checkpoint進行壓縮以節省空間呢？當我們使⽤AdamW優化器和混合精度策略來訓練一個參數量為Φ的⼤模型時，模型參數和優化器狀態的總佔⽤預計會達到14Φ的大小。72B參數的模型，會佔用接近1TB的磁盤空間。從下圖展示的Checkpoint各權重參數所佔⽤的空間量可以看出，FP32 格式的模型參數、動量1和動量2佔據了絕大部分的磁盤空間。因此，這部分是我們需要重點優化的對象。

2.2.1 AdamW優化器更新方式

大模型訓練常⽤的優化器為 Adam 優化器，其中會維護一二階優化器動量權重：momentum 和 variance，每次對模型權重的更新需要先根據一二階動量計算出 Adam更新率：

最終以下公式更新模型權重：

2.2.2 Checkpoint壓縮方案

通過大量的推導和實驗我們在PaddleNLP套件提供了兩種Checkpoint量化的⽅法，即 Checkpoint O1、 O2方法，分別對優化器權重參數進行 INT8、INT4 的優化器壓縮。

Checkpoint O1 壓縮和恢復方法

•壓縮步驟

使⽤channel-wise對稱均勻量化方法，以權重的channel維度為粒度計算每個 channel 的最大絕對值，得到這個 channel 的 scales 張量（FP32）。通過 scales 張量，以公式：

將優化器的一階動量（momentum）從 FP32 量化為 INT8 數據格式。最終 Checkpoint 中只需要存儲INT8 類型的一階動量與其對應的量化 scales。

實驗中發現，獨立量化一階動量和二階動量，恢復後計算出來的 Adam 更新率會因量化損失大幅振盪。因此本方法不直接量化二階動量，⽽採取直接量化 Adam更新率中的部分：

這⾥採取飛對稱均勻量化法將其從 FP32 量化為INT8數據格式。量化完成後，原本 FP32 數據類型的⼀一二階動量量化為 INT8 數據類型的格式保存，存儲空間變為原來的 25%。

•恢復步驟：

在Checkpoint恢復階段，⾸先將 INT8 數據類型的一階動量和 Adam 更新率根據各自的scales反量化為對應的 FP32 格式。接著需要通過FP32 的 Adam 更新率重新計算二階動量，使用公式：

Checkpoint O2 壓縮和恢復方法

• 壓縮步驟

O2 壓縮⽅式的壓縮對象和 O1 壓縮一致，即一階動量與 Adam 更新率，但壓縮方法有所不同。O2 壓縮

⽅式中，對一階動量進⾏ block-wise 的 INT4 對稱均勻量化，對 Adam更新率進行 block-wise 的 INT4

⾮對稱均勻量化，此外為了更好的保存數據，本方法通過二進製移位的⽅式將兩個 INT4 數據合併成⼀個 INT8 數據來保存，進⼀步提升Checkpoint壓縮率。

•恢復步驟：

與 O1壓縮恢復類似，但在Checkpoint恢復階段，首先會將 INT8 數據類型的一階動量和Adam 更新率合併張量根據位運算恢復成2 份 INT4 張量，再根據各自的scales反量化為對應的 FP32 格式。最後通過 FP32的 Adam 更新率重新計算二階動量。總方法示意圖如下：

對於壓縮可能帶來的精度損失，我們使⽤ Llama3-8B 在 6 個數據集上進行了共計 1000個 steps 的SFT訓 練，其中會在第 500 個 steps 時進行 Checkpoint 保存並強行讓訓練從 Checkpoint 中恢復，由下表所示，使⽤Checkpoint O1 和 Checkpoint O2 壓縮算法後，模型訓練精度是基本無損的，最終評估精度誤差在千分位以下。

從上面的表格來看 Checkpoint O1 和 Checkpoint O2算法在恢復訓練後整體對模型的精度無影響。

三、只需三行代碼，無縫升級

想讓原本1000GB 的 Checkpoint 瘦身成不到 300GB 麼，只需三行！通過以下命令可快速體驗 Unified Checkpoint 的使用。

• unified_checkpoint：總開關，控制是否使⽤ Unified Checkpoint存儲工具。

• unified_checkpoint_config：加載和存儲的優化配置：

skip_save_model_weight: 當master_weights存在時，跳過保存模型權重，與remove_master_weight 不兼容開啟。master_weight_compatible: 1. 僅當optimizer需要master_weights時，才進行加載; 2. 如果Checkpoint中不存在master_weights，則將model weight作為master_weights進行加載。remove_master_weight: 是否保存 master weight, 如果Checkpoint中不存在master_weights，則將model weight作為master_weights進行加載。async_save: 在保存Checkpoint至磁盤時做異步保存，不影響訓練過程，提高訓練效率。

• ckpt_quant_stage：是否開啟Checkpoint壓縮，支持配置為O1(INT8)壓縮和O2(INT4)壓縮⽅案。

python -u -m paddle.distributed.launch \2  --gpus "0,1,2,3,4,5,6,7" \3  run_finetune.py \4  ./config/llama/sft_argument.json \5  --unified_Checkpoint 1 \6  --unified_Checkpoint_config "async_save remove_master_weight" \7  --ckpt_quant_stage "O1"

更多參數配置選項可參考：

https://github.com/PaddlePaddle/PaddleNLP/blob/develop/docs/trainer.md

總結

PaddleNLP Unified Checkpoint是針對大模型大規模分佈式訓練設計和實現的開源存儲工具，可以自適應感知分佈式訓練策略的組合變化，還可以通過異步保存策略以及⽆損壓縮算法實現快速存儲、降低存儲空間。

在性能收益方面，Unified Checkpoint通過其獨特的異步保存機制和無損壓縮算法，實現了存儲速度的⼤幅提升和存儲空間的顯著節省。異步存儲技術使得模型訓練與保存操作能夠並行進行，大幅減少了訓練過程中的等待時間，根據測試數據，存儲耗時減少最高可達95%，這意味著訓練過程更加流暢，訓練資源的利⽤率得到了極大提高。同時，⽆損壓縮算法的應用使得Checkpoint的存儲空間需求大幅降低，最高可節省78.5%的存儲空間。

在使用效果方面，Unified Checkpoint的統一存儲格式和分佈式策略解耦設計，使得模型在不同分佈式策略間切換變得輕鬆自如，無需擔⼼因策略變化而導致的恢復訓練難題。這⼀特性極大地增強了訓練的靈活性和可擴展性，使得⽤戶能夠根據實際資源情況靈活調整訓練策略，而無需擔⼼存儲格式的兼容性問題。此外，訓練和推理模型存儲格式的無縫切換，更是為⽤戶提供了從訓練到部署的一站式解決⽅案，簡化了模型應用的複雜度，提高了開發和部署效率。

綜上所述，PaddleNLP UnifiedCheckpoint以其卓越的性能收益和便捷的使⽤效果，為大型語言模型的訓練與部署提供了強有力的⽀持。它不僅解決了大規模訓練中存儲效率和存儲空間的瓶頸問題，還通過其靈活性和可擴展性，為用戶帶來了更加⾼效、便捷的訓練體驗。

Unified Checkpoint使用文檔：

https://paddlenlp.readthedocs.io/zh/latest/llm/docs/unified_Checkpoint.html

Paddle 鏈接：

https://github.com/PaddlePaddle/Paddle

PaddleNLP 鏈接：

https://github.com/PaddlePaddle/PaddleNLP

精彩活動預告

為了讓您能夠迅速且深入地瞭解飛槳大模型套件Unified Checkpoint技術，我們在官方交流群內為大家提供了核心技術解讀影片課程。同時為了幫助大家更好地上手體驗UnifiedCheckpoint對大模型訓練流程的顯著優化，官方交流群內提供了詳細的全流程使用教程文檔，參與實戰營活動體驗任務算力全免費，成功打卡全部任務還將額外獲得AI Studio算力會員卡，官方技術人員也會提供技術答疑。機會難得，立即掃瞄下方二維碼預約吧！