大模型訓練檢查點寫入速度相比 PyTorch 加快 116 倍！

作者 | 蔡芳芳

近日，微軟 DeepSpeed 研究組釋出最新論文，提出一種名為 FastPersist 的新方法，旨在解決大模型訓練時寫檢查點十分耗時的問題，相比 PyTorch 基線，寫入速度提升超過 100 倍。

深度學習作為推動人工智能發展的關鍵技術，其模型檢查點（checkpoint）的生成對于確定訓練過程的穩定性和容錯性至關重要。然而，随着模型規模的不斷擴大，傳統的檢查點寫入方法已經無法滿足日益增長的 I/O 需求，成為制約深度學習發展的瓶頸。FastPersist 技術的提出，正是為了解決這一問題。

FastPersist 是微軟 DeepSpeed 團隊針對深度學習模型訓練中檢查點建立效率低下的問題提出的解決方案。據介紹，這項技術的核心在于通過三種創新的方法，即優化 NVMe SSDs 的使用、提高寫入并行性，以及實作檢查點操作與獨立訓練計算的重疊，顯著提升了檢查點的建立速度，降低了訓練過程中的 I/O 開銷。實驗結果表明，FastPersist 能夠在幾乎不影響訓練性能的前提下，實作高達 116 倍的檢查點寫入速度提升。這項技術的提出，不僅解決了大規模深度學習訓練中的一個關鍵問題，也為未來深度學習模型的進一步發展提供了強有力的技術支援。

AI 前線進一步了解到， 在微軟很多重要的大模型訓練中，由于工作負載高度密集，經常出現 GPU error，是以需要很高頻地寫檢查點操作，而這些大模型訓練其實都在使用 FastPersist 這套系統。

論文連結：https://arxiv.org/pdf/2406.13768

現狀和問題

深度學習作為人工智能領域的一個重要分支，近年來在圖像識别、自然語言處理、推薦系統等多個領域取得了突破性進展。随着研究的深入，深度學習模型的規模也在不斷擴大，從早期的百萬級參數模型發展到現在的百億甚至千億級參數的超大型模型。模型規模的增長帶來了更強的表征能力和更高的準确率，但同時也帶來了計算複雜度的提升和存儲需求的增加。特别是模型參數、梯度資訊以及中間特征圖等資料的存儲，對存儲系統的 I/O 性能提出了更高的要求。

盡管計算性能的提升可以通過硬體加速和算法優化來實作，但 I/O 性能的提升卻受到了傳統儲存設備和系統的限制。特别是在模型訓練過程中，檢查點的生成是一個不可或缺的步驟，用于儲存模型在特定疊代步驟的狀态，以便在發生故障時能夠從最近的檢查點恢複訓練，進而避免重複計算。然而，檢查點的生成和儲存是一個資源密集型的操作，涉及到大量的資料寫入。在大規模訓練中，模型參數和中間資料的體積巨大，檢查點的生成和儲存需要消耗大量的 I/O 帶寬和時間，這不僅增加了訓練的總體時間，也可能導緻 I/O 系統的飽和，影響其他訓練操作的執行。是以，提高檢查點建立的效率，成為提升深度學習模型訓練性能的關鍵。

目前深度學習架構中的檢查點生成機制，大多數基于傳統的檔案 I/O 操作，這些操作并沒有充分利用現代儲存設備，如 NVMe SSDs 的高性能特性。這導緻了在大規模訓練場景下，檢查點寫入成為制約整體性能的瓶頸。此外，由于檢查點寫入操作與模型訓練的其他計算任務之間存在資料依賴性，傳統的檢查點生成方法無法實作與訓練過程的完全解耦，進一步限制了檢查點生成的效率。

為了解決 I/O 瓶頸問題，研究者和工程師們提出了多種解決方案，如使用更快的存儲媒體、優化檔案系統、改進資料寫入政策等。但是，這些解決方案往往存在一定的局限性。例如，簡單地更換更快的存儲媒體雖然可以提高 I/O 性能，但成本較高，且在大規模并發寫入時仍可能遇到瓶頸。優化檔案系統和資料寫入政策可以在一定程度上提高效率，但往往需要對現有的深度學習架構和訓練流程進行較大的改動，相容性和通用性有待提高。

針對上述問題，微軟 DeepSpeed 團隊提出了 FastPersist 技術。

FastPersist 技術方案

FastPersist 通過深入分析深度學習訓練過程中的 I/O 需求和特點，結合現代儲存設備的特性，提出了一種全新的檢查點生成和儲存方法。主要通過以下三個方面來提升檢查點建立的效率：

1. NVMe 儲存設備的優化利用

FastPersist 針對 NVMe SSDs 的高性能特性進行了優化。通過使用專為 NVMe 設計的 I/O 庫，如 libaio 和 io_uring，FastPersist 能夠更高效地管理資料在 GPU 和 SSD 之間的傳輸，進而顯著提高了單節點上的檢查點寫入速度。

FastPersist 還采用了雙緩沖技術來進一步提高寫入效率。在雙緩沖機制中，當一個緩沖區的資料正在寫入 SSD 時，另一個緩沖區可以同時從 GPU 記憶體中預取資料，這樣就能實作資料寫入和資料預取的流水線操作，減少了等待時間，提高了整體的寫入性能。

另外， FastPersist 針對 NVMe SSDs 的特性，對資料塊的大小和對齊進行了優化。通過調整資料塊的大小，使其比對 SSD 的頁面大小，可以減少寫入操作的數量，提高寫入效率。同時，通過對齊資料塊到合适的邊界，可以避免額外的拷貝操作，進一步提高性能。

2. 寫入并行性的實作

在深度學習模型訓練中，特别是在大規模分布式訓練環境中，資料并行（Data Parallelism）是一種常見的訓練政策。在資料并行訓練中，模型被複制到多個訓練節點上，每個節點處理不同的資料子集。這種訓練方式可以顯著提高計算資源的使用率，加快模型的訓練速度。然而，如果檢查點的寫入操作仍然集中在單個節點上執行，那麼 I/O 操作就可能成為限制整體性能的瓶頸。

FastPersist 技術通過實作檢查點寫入的并行性，解決了這一問題。在 FastPersist 中，檢查點的寫入操作被分布到所有參與訓練的節點上，每個節點隻負責寫入其對應的模型部分。這樣，寫入操作就可以同時在多個節點上執行，進而顯著提高了整體的寫入速度。

為了實作高效的寫入并行性，FastPersist 采用了以下幾個關鍵政策：

1. 資料分片：FastPersist 将檢查點資料均勻地分割成多個片段，每個訓練節點隻負責寫入其配置設定到的資料片段。這種分片政策確定了寫入負載在所有節點上的均衡配置設定。
2. 無通信寫入：在 FastPersist 中，每個節點獨立地完成其檢查點資料片段的寫入，無需與其他節點進行通信或協調。這種設計減少了節點間通信的開銷，提高了寫入操作的效率。
3. 動态負載平衡：FastPersist 能夠根據節點的計算能力和存儲性能動态調整資料片段的大小，確定所有節點的寫入負載保持均衡。這種動态調整機制可以适應不同的硬體環境和訓練配置。
4. 容錯和恢複：在分布式訓練環境中，節點的故障是不可避免的。FastPersist 通過在寫入操作中實作容錯機制，確定即使部分節點發生故障，也不會影響檢查點的完整性和訓練的連續性。

3. 操作重疊的政策

在深度學習模型訓練中，檢查點的生成通常需要在每個訓練疊代後執行，以確定模型狀态的持久化。然而，如果每次疊代後都進行完整的檢查點寫入操作，那麼這些操作可能會占用大量的計算資源，影響模型訓練的速度。為了解決這一問題，FastPersist 采用了操作重疊的政策，将檢查點的寫入操作與模型訓練的其他計算任務并行執行。

操作重疊的核心思想是利用深度學習訓練中的計算特性，将檢查點寫入操作與模型的前向傳播和後向傳播操作重疊。由于前向傳播和後向傳播操作通常占據了模型訓練的大部分時間，通過将檢查點寫入操作與這些操作并行化，可以有效地隐藏 I/O 操作的延遲，提高整體的訓練效率。

FastPersist 實作操作重疊的具體政策包括：

1. 異步寫入：FastPersist 采用異步寫入機制，使得檢查點的寫入操作不會阻塞計算操作的執行。在每個訓練疊代的優化器步驟之後，FastPersist 會啟動檢查點的異步寫入過程，而計算線程可以繼續執行下一個疊代的前向傳播和後向傳播。
2. 雙線程模型：FastPersist 引入了一個輔助線程專門負責檢查點的寫入操作。主線程負責執行模型的計算任務，而輔助線程在主線程的協調下執行檢查點的寫入。這種雙線程模型確定了計算和 I/O 操作的并行執行，減少了互相之間的幹擾。
3. 資料局部性優化：FastPersist 通過優化資料的存儲和通路模式，提高了資料在 GPU 和 CPU 之間的傳輸效率。通過利用資料的局部性原理，FastPersist 減少了不必要的資料移動，降低了 I/O 操作的延遲。
4. 依賴性管理：在操作重疊的過程中，FastPersist 通過精确管理計算任務和檢查點寫入操作之間的資料依賴性，確定了檢查點的一緻性和完整性。即使在發生故障的情況下，FastPersist 也能夠保證從最近的檢查點正确恢複。

通過精心設計的操作排程政策，FastPersist 實作了檢查點寫入操作與模型訓練的其他計算任務的重疊執行，進而在不增加額外計算負擔的情況下，規避檢查點的寫入延遲。

效果評估

研究團隊對 FastPersist 的性能表現進行多場景、多元度的評估。為了驗證 NVMe 優化和并行優化在減少檢查點延遲方面的有效性，團隊使用單 GPU 和多節點環境的微基準測試，對檢查點寫入的吞吐量做了測試；并使用真實世界的密集和稀疏深度學習模型，評估了新方法相比基線（baseline）對訓練性能的加速效果。

在微基準測試中，FastPersist 在單 GPU 和多節點環境下，相比于基線的 torch.save()方法，檢查點寫入速度顯著提升。

在真實世界的深度學習模型訓練測試中，FastPersist 在不同的模型規模和資料并行度下，均能夠實作高速的檢查點建立，且引入的開銷極小。下圖顯示，在 128 個 V100 GPU 上，FastPersist 實作的加速比從 gpt3-13B 的 28 倍到 gpt3-0.7B 的 116 倍不等。這些改進證明了 FastPersist 技術方案在 NVMe 優化和并行優化方面的有效性。

圖：FastPersist 應用于 GPT-3 密集模型訓練的效果

FastPersist 在大規模訓練場景下的性能尤為重要。實驗結果表明，即使在數千個 GPU 上進行訓練，FastPersist 也能夠保持檢查點建立的低開銷，并且随着資料并行度的增加，FastPersist 的效率提升更加明顯。

鑒于 GPU 硬體的限制，研究團隊通過預測高達 128 的資料并行度（即 6.7B 模型使用 1024 個 GPU，13B 模型使用 2048 個 GPU）來模拟像 GPT-3 6.7B 和 13B 這樣的大型密集模型的性能表現。下圖顯示了 FastPersist 相對于基線的預計訓練加速比，其中藍色/橙色條代表 6.7B/13B 模型。當擴充到數千個 GPU 時，FastPersist 的檢查點開銷基本保持一緻（小于 2%的訓練計算時間），而基線的檢查點開銷則與資料并行度成比例增長。對于 6.7B 和 13B 模型，FastPersist 分别實作了高達 10.2 倍和 3.6 倍的訓練加速。

圖：資料并行度≤128 的訓練加速效果預測

另外如上圖中灰色條所示，如果放棄流水線并行（PP），并在一個資料并行組中完全采用 16 個 GPU 的張量并行（TP）設定，與标準 TP 和 PP 結合的模型分割（即圖中的橙色條）相比，FastPersist 可以做到更高的基線加速比，實作高達 11.3 倍的訓練加速。

原文連結：大模型訓練檢查點寫入速度相比 PyTorch 加快 116 倍！微軟提出FastPersist 新方法_微軟_蔡芳芳_InfoQ精選文章