伴隨大模型迭代速度越來越快，訓練叢集規模越來越大，高頻率的軟硬體故障已經成為阻礙訓練效率進一步提高的痛點，檢查點（Checkpoint）系統在訓練過程中負責狀態的儲存和恢復，已經成為克服訓練故障、保障訓練進度和提高訓練效率的關鍵。
近日，位元組跳動豆包大模型團隊與香港大學聯合提出了 ByteCheckpoint。這是一個 PyTorch 原生，相容多個訓練框架，支援 Checkpoint 的高效讀寫和自動重新切分的大模型 Checkpointing 系統，相比現有方法有顯著效能提升和易用性優勢。本文介紹了大模型訓練提效中 Checkpoint 方向面臨的挑戰，總結 ByteCheckpoint 的解決思路、系統設計、I/O 效能最佳化技術，以及在儲存效能和讀取效能測試的實驗結果。
Meta 官方最近披露了在 16384 塊 H100 80GB 訓練叢集上進行 Llama3 405B 訓練的故障率 —— 短短 54 天，發生 419 次中斷，平均每三小時崩潰一次，引來不少從業者關注。
正如業內一句常言，大型訓練系統唯一確定的，便是軟硬體故障。隨著訓練規模與模型大小的日益增長，克服軟硬體故障，提高訓練效率成為大模型迭代的重要影響要素。
Checkpoint 已成為訓練提效關鍵。在 Llama 訓練報告中，技術團隊提到，為了對抗高故障率，需要在訓練過程中頻繁地進行 Checkpoint ，儲存訓練中的模型、最佳化器、資料讀取器狀態，減少訓練進度損失。
位元組跳動豆包大模型團隊與港大近期公開了成果 —— ByteCheckpoint ，一個 PyTorch 原生，相容多個訓練框架，支援 Checkpoint 的高效讀寫和自動重新切分的大模型 Checkpointing 系統。
與基線方法相比，ByteCheckpoint 在 Checkpoint 儲存上效能提升高達 529.22 倍，在載入上，效能提升高達 3.51 倍。極簡的使用者介面和 Checkpoint 自動重新切分功能，顯著降低了使用者上手和使用成本，提高了系統的易用性。
- ByteCheckpoint: A Unified Checkpointing System for LLM Development
- 論文連結：https://team.doubao.com/zh/publication/bytecheckpoint-a-unified-checkpointing-system-for-llm-development?view_from=research
Checkpoint 技術在大模型訓練中的技術挑戰 當前 Checkpoint 相關技術在支援大模型訓練提效中，共面臨四個方面挑戰：- 現有系統設計存在缺陷，顯著增加訓練額外 I/O 開銷
在訓練工業級別的大語言模型 (LLM) 的過程中，訓練狀態需要透過檢查點技術 ( Checkpointing ) 進行儲存和持久化。通常情況下，一個 Checkpoint 包括 5 個部分 (模型，最佳化器，資料讀取器，隨機數和使用者自定義配置)。這一過程往往會給訓練帶來分鐘級別的阻塞，嚴重影響訓練效率。在使用遠端持久化儲存系統的大規模訓練場景下，現有的 Checkpointing 系統沒有充分利用 Checkpoint 儲存過程中 GPU 到 CPU 記憶體複製 ( D2H 複製)，序列化，本地存檔，上傳到儲存系統各個階段的執行獨立性。此外，不同訓練程序共同分擔 Checkpoint 存取任務的並行處理潛力也沒有被充分發掘。這些系統設計上的不足增加了 Checkpoint 訓練帶來的額外 I/O 開銷。- Checkpoint 重新切分困難，手動切分指令碼開發維護開銷過高
在 LLM 的不同訓練階段 (預訓練到 SFT 或者 RLHF ) 以及不同任務 (從訓練任務拉取不同階段的 Checkpoint 進行執行自動評估) 之間進行 Checkpoint 遷移時，通常需要對儲存在持久化儲存系統中的 Checkpoint 進行重新切分 ( Checkpoint Resharding ) ，以適應下游任務的新並行度配置以及可用 GPU 資源的配額。現有 Checkpointing 系統 [1, 2, 3, 4] 都假設儲存和載入時，並行度配置和 GPU 資源保持不變，無法處理 Checkpoint 重新切分的需求。工業界目前常見的解決辦法是 —— 為不同模型定製 Checkpoint 合併或者重新切分指令碼。這種方法帶來了大量開發與維護開銷，可擴充套件性較差。- 不同的訓練框架 Checkpoint 模組割裂，為 Checkpoint 統一管理和效能最佳化帶來挑戰
在工業界的訓練平臺上，工程師與科學家往往會根據任務特性，選擇合適框架 (Megatron-LM [5], FSDP [6], DeepSpeed [7], veScale [8, 9]) 進行訓練，並儲存 Checkpoint 到儲存系統。然而，這些不同的訓練框架都具有自己獨立的 Checkpoint 格式以及讀寫模組。不同訓練框架的 Checkpoint 模組設計不盡相同，為底層系統進行統一的 Checkpoint 管理以及效能最佳化帶來了挑戰。從訓練系統的使用者（ AI 研究科學家或工程師）的角度出發，使用者使用分散式訓練系統時，在 Checkpoint 方向往往會被三個問題困擾：1）如何高效地儲存 Checkpoint ，在不影響訓練效率的情況下儲存 Checkpoint。2）如何重新切分 Checkpoint ，對於在一個並行度下儲存的 Checkpoint ，根據新的並行度正確讀入。3）如何把訓練得到的產物上傳到雲端儲存系統上（ HDFS，S3 等），手動管理多個儲存系統，對使用者來說學習和使用成本較高。針對上述問題，位元組跳動豆包大模型團隊和香港大學吳川教授實驗室聯合推出了 ByteCheckpoint 。ByteCheckpoint 是一個多訓練框架統一，支援多儲存後端，具備自動 Checkpoint 重新切分能力的高效能分散式 Checkpointing 系統。ByteCheckpoint 提供了簡單易用的使用者介面 ，實現了大量 I/O 效能最佳化技術提高了儲存和讀取 Checkpoint 效能，並支援 Checkpoint 在不同並行度配置的任務中的靈活遷移。ByteCheckpoint 採用了後設資料 / 張量資料分離的儲存架構，實現了 Checkpoint 管理與訓練框架和並行度的解耦合。不同訓練框架中的模型以及最佳化器的張量切片 ( Tensor Shard) 儲存在 storage 檔案中，元資訊 (TensorMeta, ShardMeta, ByteMeta) 儲存到全域性唯一的 metadata 檔案中。
當使用不同的並行度配置讀取 Checkpoint 時，如下圖所示，每個訓練程序只需要根據當前的並行度設定查詢元資訊，便能夠獲取程序所需要張量的儲存位置，再根據位置直接讀取，實現自動 Checkpoint 重新切分。
不同訓練框架在執行時，往往會把模型或者最佳化器中張量的形狀攤平 ( Flatten ) 成一維，從而提高集合通訊效能。這種攤平操作給 Checkpoint 儲存帶來了不規則張量切分 (Irregular Tensor Sharding) 的挑戰。如下圖所示，在 Megatron-LM (由 NVIDIA 研發的分散式大模型訓練框架) 和 veScale (由位元組跳動研發的 PyTorch 原生分散式大模型訓練框架) 中，模型引數對應的最佳化器狀態會被展平為一維後合併，再根據資料並行度切分。這導致張量被不規則地切分到不同程序之中，張量切片的元資訊無法使用偏移量和長度元組來表示，給儲存和讀取帶來困難。
不規則張量切分的問題在 FSDP 框架中也同樣存在。為消除不規則切分的張量切片 ，FSDP 框架在儲存 Checkpoint 之前會在所有程序上對一維張量切片進行 all-gather 集合通訊以及 D2H 複製操作，以獲取完整不規則切分的張量。這種方案帶來了極大的通訊和頻繁的 GPU-CPU 同步開銷，嚴重影響了 Checkpoint 儲存的效能。針對這個問題，ByteCheckpoint 提出了非同步張量合併 (Asynchronous Tensor Merging) 技術。ByteCheckpoint 首先找出不同程序中被不規則切分的張量，之後採用非同步的 P2P 通訊，把這些不規則的張量分配到不同程序上進行合併。所有針對這些不規則張量的 P2P 通訊等待（Wait) 以及張量 D2H 複製操作被推遲到他們即將進入序列化階段的時候，從而消除了頻繁的同步開銷，也增加了通訊與其他 Checkpoint 儲存流程的執行重疊度。下圖展示了 ByteCheckpoint 的系統架構： API 層為不同訓練框架提供了簡單，易用且統一的讀取和寫入 ( Save ）和讀取（ Load ）介面。Planner 層會根據存取物件為不同訓練程序生成存取方案，交由 Execution 層執行實際的 I/O 任務。Execution 層執行 I/O 任務並與 Storage 層進行互動，利用各種 I/O 最佳化技術進行高效能的 Checkpoint 存取。Storage 層管理不同的儲存後端，並在 I/O 任務過程中根據不同儲存後端進行相應的最佳化。分層設計增強了系統的可擴充套件性，以便未來支援更多的訓練框架和儲存後端。
ByteCheckpoint 的 API 用例如下：
ByteCheckpoint 提供了極簡 API ，降低了使用者上手的成本。使用者在儲存和讀取 Checkpoint 時，只需要呼叫儲存和載入函式，傳入需要儲存和讀取的內容，檔案系統路徑和各種效能最佳化選項。如下圖所示，ByteCheckpoint 設計了全非同步的儲存流水線（Save Pipeline），將 Checkpoint 儲存的不同階段（P2P 張量傳輸，D2H 複製，序列化，儲存本地和上傳檔案系統）進行拆分，實現高效的流水線執行。
在 D2H 複製過程，ByteCheckpoint 採用固定記憶體池（ Pinned Memory Pool ），減少了記憶體反覆分配的時間開銷。除此之外，為了降低高頻儲存場景中因為同步等待固定記憶體池回收而帶來的額外時間開銷，ByteCheckpoint 在固定記憶體池的基礎上加入了 Ping-Pong buffering 的機制。兩個獨立的記憶體池交替扮演著讀寫 buffer 的角色，與 GPU 和執行後續 I/O 操作的 I/O workers 進行互動，進一步提升儲存效率。
在資料並行 ( Data-Parallel or DP ) 訓練中，模型在不同的資料並行程序組（ DP Group ）之間是冗餘的， ByteCheckpoint 採用了負載均衡演算法把冗餘的模型張量均勻分配到不同程序組中進行儲存，有效地提高了 Checkpoint 儲存效率。如圖所示，在改變並行度讀取 Checkpoint 時，新的訓練程序可能只需要從原來的張量切片中讀取其中的一部分。ByteCheckpoint 採用按需部分檔案讀取（ Partial File Reading ） 技術，直接從遠端儲存中讀取需要的檔案片段，避免下載和讀取不必要的資料。
在資料並行 (Data-Parallel or DP) 訓練中，模型在不同的資料並行程序組（DP Group）之間是冗餘的，不同程序組會重複讀取同一個張量切片。在大規模訓練的場景下，不同程序組同時發給遠端持久化儲存系統 （比如 HDFS ）大量請求，會給儲存系統帶來巨大壓力。為了消除重複資料讀取，減少訓練程序發給 HDFS 的請求，最佳化載入的效能，ByteCheckpoint 把相同的張量切片讀取任務均勻分配到不同程序上，並在對遠端檔案進行讀取的同時，利用 GPU 之間閒置的頻寬進行張量切片傳輸。
團隊使用 DenseGPT 與 SparseGPT 模型 (基於 GPT-3 [10] 結構實現)，在不同模型引數量，不同訓練框架和不同規模的訓練任務中評估了 ByteCheckpoint 的 Checkpoint 存取正確性、儲存效能和讀取效能。更多實驗配置和正確性測試細節請移步完整論文。
在儲存效能測試中，團隊比較了不同模型規模和訓練框架，在訓練過程中每 50 或者 100 步存一次 Checkpoint ， Bytecheckpoint 和基線（ Baseline ）方法給訓練帶來的總的阻塞時間 ( Checkpoint stalls )。得益於對寫入效能的深度最佳化，ByteCheckpoint 在各類實驗場景中均取得了很高的表現，在 576 卡 SparseGPT 110B - Megatron-LM 訓練任務中相比基線儲存方法取得了 66.65~74.55 倍的效能提升，在 256 卡 DenseGPT 10B - FSDP 訓練任務中甚至能達到 529.22 倍的效能提升。
在讀取效能測試中，團隊比較不同方法根據下游任務並行度讀取 Checkpoint 的載入時間。ByteCheckpoint 相比基線方法取得了 1.55 ～ 3.37 倍的效能提升。團隊觀察到 ByteCheckpoint 相對於 Megatron-LM 基線方法的效能提升更為顯著。這是因為 Megatron-LM 在讀取 Checkpoint 到新的並行度配置之前，需要執行離線的指令碼對分散式 Checkpoint 進行重新分片。相比之下，ByteCheckpoint 能夠直接進行自動 Checkpoint 重新切分，無需執行離線指令碼，高效完成讀取。
最後，關於 ByteCheckpoint 的未來規劃，團隊希望從兩個方面著手：其一，實現支援超大規模 GPU 叢集訓練任務高效 Checkpointing 的長遠目標。其二，實現大模型訓練全生命週期的 Checkpoint 管理，支援全場景的 Checkpoint ，從預訓練（Pre-Training)，到監督微調（ SFT ），再到強化學習（ RLHF ）和評估 (Evaluation) 等場景。位元組跳動豆包大模型團隊成立於 2023 年，致力於開發業界最先進的 AI 大模型技術，成為世界一流的研究團隊，為科技和社會發展作出貢獻。目前，團隊正在持續吸引優秀人才加入，硬核、開放且充滿創新精神是團隊氛圍關鍵詞，團隊致力於創造一個積極向上的工作環境，鼓勵團隊成員不斷學習和成長，不畏挑戰，追求卓越。希望與具備創新精神、責任心的技術人才一起，推進大模型訓練提效工作取得更多進展和成果。[1] Mohan, Jayashree, Amar Phanishayee, and Vijay Chidambaram. "{CheckFreq}: Frequent,{Fine-Grained}{DNN} Checkpointing." 19th USENIX Conference on File and Storage Technologies (FAST 21). 2021.[2] Eisenman, Assaf, et al. "{Check-N-Run}: A Checkpointing system for training deep learning recommendation models." 19th USENIX Symposium on Networked Systems Design and Implementation (NSDI 22). 2022.[3] Wang, Zhuang, et al. "Gemini: Fast failure recovery in 分散式 training with in-memory Checkpoints." Proceedings of the 29th Symposium on Operating Systems Principles. 2023.[4] Gupta, Tanmaey, et al. "Just-In-Time Checkpointing: Low Cost Error Recovery from Deep Learning Training Failures." Proceedings of the Nineteenth European Conference on Computer Systems. 2024.[5] Shoeybi, Mohammad, et al. "Megatron-lm: Training multi-billion parameter language models using model parallelism." arXiv preprint arXiv:1909.08053 (2019).[6] Zhao, Yanli, et al. "Pytorch fsdp: experiences on scaling fully sharded data parallel." arXiv preprint arXiv:2304.11277 (2023).[7] Rasley, Jeff, et al. "Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters." Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020.[8] Jiang, Ziheng, et al. "{MegaScale}: Scaling large language model training to more than 10,000 {GPUs}." 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24). 2024.[9] veScale: A PyTorch Native LLM Training Framework https://github.com/volcengine/veScale[10] Brown, Tom, et al. "Language models are few-shot learners." Advances in neural information processing systems 33 (2020): 1877-1901.