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

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

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

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

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

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

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

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

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

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

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

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

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

為了實現高效的寫入並行性，FastPersist 採用了以下幾個關鍵策略：

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

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

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 在大規模訓練場景下的效能尤為重要。實驗結果表明，即使在數千個 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 倍的訓練加速。

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