大模型儲存選型 & JuiceFS 在關鍵環節效能詳解

JuiceFS發表於2024-10-11

從去年開始,LLM大語言模型領域發展迅速、如 LLaMA、ChatGLM、Baichuan、Qwen 和 yi-model 等基礎模型(Foundation Models)的數量顯著增加。眾多企業也開始基於這些基礎模型做 post-training 的相關工作,以開發特定垂直領域的模型實現應用落地。

AI 模型的引數規模呈指數級增長,出現了越來越多的千億甚至萬億引數的通用大模型。例如,最新的 llama3 模型就提供了 405B、70B 和 8B 三個版本,模型引數量不斷重新整理,給企業儲存的成本與規劃帶來很大挑戰。除了模型規模的增長,大模型開發的複雜流程和高效的資料流轉也對儲存系統提出了更高的要求,這直接影響到整個流程的效率。

另外,大規模訓練和推理對於GPU算力的需求,導致算力的供給越來越多的轉變為多雲、多region 的模式,在這種情況下,又給使用者帶來了資料集、模型在多區域間進行分發、同步以及一致性管理的挑戰。

在本文中,我們將簡要介紹我們對大模型儲存選型的思考、開發中不同階段的工作負載需求,包括成本、效能和功能等多個維度,並探討如何透過 JuiceFS 幫助使用者最佳化這些環節。

01 大模型儲存選型思考

針對資料量大和資料流轉複雜的挑戰,通常採用的解決方案是建立統一的資料管理也就是資料湖。這種方法透過統一名稱空間,在後端儲存不同型別的資料,同時前端提供多樣化的資料訪問方式,有效減少了複雜資料管道中資料流轉的難度和成本。

儘管理想很完美,但現實中如何設計一個統一的技術棧或產品來實現這些目標頗具挑戰。從支援前端資料處理框架的角度來看,全面相容 POSIX 協議的檔案系統目前看來是最佳方案。

此外,檔案系統還需支援龐大的檔案規模,既要能支援數百億檔案的儲存,又要保證在此規模下的檔案訪問效能;這就需要檔案系統具備高效的後設資料管理能力和多級快取等高效能實現。

還包括對其他容器平臺儲存編排的支援、靈活的橫向擴充套件能力以及多雲多資料中心的資料同步支援等,這些都是選擇資料湖產品時需要考慮的關鍵因素。

以下表格中列舉了市場上一些常見的檔案系統產品,並進行了多維度的橫向比較。希望能為大家選型時提供參考。

JuiceFS 使用者在選型過程中,常詢問的產品包括 CephFS、Lustre 和 Alluxio ,我們將重點梳理這 3 個產品的差異。

我們從產品定位角度來比較這幾款產品。

在對資料處理框架的支援方面,需要更好的 POSIX 協議相容性。

  • JuiceFS、CephFS 和 Lustre 都是檔案系統,因此它們均提供較全面的 POSIX 支援並確保資料在不同環節的強一致性;
  • Alluxio 主要定位為資料編排和加速工具,只部分支援 POSIX 且不保證資料的強一致性,因此適用的使用場景有所差異。

在支援海量小檔案方面,挑戰主要體現在檔案後設資料的管理能力和效能上。

  • CephFS 和 Lustre 利用磁碟儲存後設資料,並透過記憶體快取進行效能加速;檔案規模達到一定量級後,後設資料的管理和使用成本會明顯增加;
  • Alluxio 提供基於 RocksDB 和基於堆記憶體的兩種後設資料儲存方案,使用者可以根據自己的需求選擇適合的方案;
  • JuiceFS 使用物件儲存作為後端,該儲存方式具有扁平名稱空間、高可擴充套件性和低成本等特點,非常適合儲存海量檔案。社群版的後設資料引擎採用開放的外掛架構設計,使用者可以根據檔案規模、效能需求、使用習慣等從多種主流的開源資料庫中進行選擇來作為後設資料儲存, 而企業版的則是自研了基於全記憶體實現的後設資料引擎服務,針對檔案目錄樹管理,從後設資料訪問效能、橫向擴充套件能力幾個方面進行定製最佳化。

在多雲資料管理方面,需要對使用者透明的檔案系統級別複製能力。

  • CephFS 和 Lustre本身不具備檔案系統級別複製功能,需要透過儲存底座的複製功能或者第三方工具來實現,在資料一致性上難以保障,也會增加額外的管理成本。
  • JuiceFS 和Alluxio 都提供了對使用者透明的多雲、多 region 的資料分發能力。除此以外,JuiceFS 還提供了在映象站點檔案系統可寫的能力;

總結一下 JuiceFS 在資料湖儲存中的給使用者帶來的收益:

  • 統一名稱空間提供多協議訪問;
  • 完整的 POSIX 協議相容性;百億級別檔案規模的支援;
  • 資料強一致性;
  • 高併發共享訪問能力;
  • 不同負載的效能最佳化。
  • 多雲、多region的資料分發和訪問。

02 JuiceFS 在大模型場景的實踐

環節1:資料集載入

大模型訓練的資料集載入有兩個關鍵點:資料集需要被反覆多輪遍歷,即多個 epoch,並且在每個 epoch 開始之前資料集需要被隨機打亂。這樣做是為了確保模型能夠充分學習資料中的模式,並避免過擬合。接著這些被隨機打散的資料集按批次載入到 GPU 視訊記憶體中進行處理。

由於這些資料處理主要依賴 CPU 資源,因此 GPU 在這一階段往往處於空閒。使用者需要提升資料集載入的效率,如使用快取,從而減少 GPU 的閒置時間。資料集通常以結構化大檔案(如 FFRecord、TFRecord、Arrow)或海量小檔案,如未打包的文字、圖片、音訊等形式。

資料集的隨機打散過程需要隨機讀取資料檔案,結合資料集檔案的格式,資料集載入流程對於儲存的 I/O 模型是大、小檔案隨機讀。隨機讀 I/O 需要檔案系統具備高 IOPS 和低 I/O 延遲,尤其對於海量小檔案隨機讀的 I/O 處理,對檔案系統後設資料引擎的效能提出了更高要求。另外資料集在一次訓練過程被反覆讀取,那資料如果都能被快取在高效能儲存介質中,實現高快取命中率,就可以最大程度的提升資料讀取的效能。

JuiceFS 的設計目標是幫助使用者在效能和成本之間找到最佳平衡。它透過利用物件儲存作為後端資料持久化手段,並結合多級快取加速架構,為資料集載入提供了高 IOPS 和低 IO 延遲,同時保持成本效益。下面列出了使用 JuiceFS 進行小 IO 隨機讀操作的一些效能指標:

  • 命中本地快取時,延遲在 0.05 到 0.1 毫秒之間;
  • 命中分散式快取(企業版特性)時,延遲在 0.1 到 0.2 毫秒之間;
  • 直接從物件儲存讀取時,延遲超過 30 到 50毫秒。

使用 libaio 可以獲得更高的 IOPS,但這需要整合特定框架擴充套件以支援 libaio 介面。另外,JuiceFS 企業版提供自研的全記憶體實現的後設資料引擎,可以在管理海量檔案規模的情況下,保證後設資料請求的平均處理時間維持在 100 微秒量級,在海量小檔案的資料集載入場景中,也得到很好的效能。另外透過後設資料引擎多分割槽的架構實現,提供了根據檔案規模動態線性的擴充套件能力。

下面附上一組使用者的大檔案隨機讀測試資料,幫助大家理解 JuiceFS 的效能。測試資料為單個 1.8TB 檔案。

環節2:訓練環節 checkpoint 儲存

對訓練環節的 checkpoint 檔案進行儲存,是為了在訓練中斷時可以從最近的 checkpoint 進行訓練恢復,避免從頭開始,從而節約時間和計算資源。如果採用同步寫 checkpoint,那麼在儲存 checkpoint 的時間視窗內,GPU 會全部等待。通常,checkpoint 寫入都是大檔案的順序寫。要儘量縮短這個過程的時間,就需要檔案系統提供高效能的寫吞吐能力

JuiceFS 透過對接物件儲存作為資料持久化層,其吞吐上限受限於物件儲存的寫頻寬、專線頻寬以及 JuiceFS 客戶端所在節點的網路卡頻寬。JuiceFS 採用分塊儲存設計,透過增加對物件儲存的併發,可以充分利用物件儲存頻寬,來提升大檔案順序寫的吞吐能力。

此外,JuiceFS 還提供寫快取功能。若後端物件儲存存在效能瓶頸,可考慮開啟 writeback 寫快取,利用 GPU 節點的 NVMe 高速儲存介質作為本地快取,首先將 checkpoint 檔案寫入本地,然後再非同步上傳到物件儲存,這有助於減少寫入延遲。這裡需要注意,開啟 writeback 寫快取,可能會在一些場景中導致 checkpoint 檔案不一致,無法正常載入恢復訓練,針對這種情況就需要載入上一個視窗的 checkpoint 檔案進行恢復,從而重跑一些訓練 epcoh。

根據使用者實際應用的驗證,在透過 torch.save 向 JuiceFS 檔案系統儲存 checkpoint 的場景,每個 GPU 單卡啟動的一個寫 checkpoint 程序, 可以達到 1.7GB/s 或更高的寫吞吐,能夠滿足該場景的效能要求

環節3:訓練和推理環節的模型載入

典型的模型載入場景包括:

  • 訓練恢復載入 Checkpoint 檔案場景:在訓練恢復階段,採用並行訓練的節點僅載入其 rank 的分片 checkpoint 檔案,所有節點需要載入的總檔案大小就是所有分片 checkpoint 大小的總和,載入 checkpoint 檔案的效率決定了 GPU 等待訓練恢復的時長,直接影響 GPU 的利用率。

  • 推理服務載入模型檔案場景: 將訓練好的模型部署到推理服務中,一般需要在推理節點啟動時,載入完整的模型檔案。當推理節點的數量很多時,如執行上千個推理節點,同時載入模型,每個節點都需要從檔案儲存上去讀取一個完整的模型檔案,這時會產生非常大的讀吞吐,如果網路吞吐成為瓶頸,模型的載入效率就會很慢。如果檔案儲存是採用公有云上的物件儲存,那物件儲存的頻寬、專線頻寬等都很容易成為瓶頸,而且每個節點都從物件儲存 get 完整的模型檔案,也會產生大量額外的物件儲存頻寬和呼叫成本。

值得注意的是,在模型載入環節中,模型檔案的不同格式,對儲存 I/O 模型的需求不一樣。模型檔案可能會採用 torch.save 儲存的 .pt、.bin 檔案,或使用 safetensors 庫儲存的 .safetensors 檔案。其中 safetensors 檔案較為特殊,載入時會採用 mmap 方式,對儲存的 I/O 模型為大檔案的隨機讀;而其他格式檔案的載入則是大檔案的順序讀。針對這兩類不同格式的模型載入的需求,來看一下 JuiceFS 的方案。

JuiceFS 對 safetensors 格式模型檔案載入(大檔案隨機讀)的最佳化實踐:包括提前將這些檔案預熱到快取中,以獲得高檔案讀取 IOPS 和低延遲的效能。最終效能將受限於快取介質的 IOPS 能力和JuiceFS 客戶端的網路環境。此外,為減少預讀可能導致的讀放大問題,考慮關閉 prefetch 預取功能。若需進一步最佳化 safetensors 檔案的載入效能,可以在模型載入前預讀檔案到核心的 pagecache 中,這一步驟能顯著提高 safetensors 檔案的載入速度,效能相較於從 JuiceFS 快取中讀取可提高一個數量級。

JuiceFS 對其他格式模型檔案載入(大檔案順序讀)的最佳化策略:JuiceFS 透過預讀功能來提升大檔案順序讀的效能。此功能可以預測並提前載入未來可能請求的資料塊到記憶體中,透過合理配置預讀視窗大小並增加訪問物件儲存的併發數,充分利用物件儲存頻寬。最終讀吞吐效能的上限主要受限於物件儲存的頻寬或到物件儲存專線的頻寬。此外,預熱 checkpoint 或模型檔案到快取中可以進一步提高效能,例如 8 卡單機載入 checkpoint 的情況下,在網路卡頻寬足夠的情況下,熱讀的吞吐量可達到 10GB/s 以上,足以應對使用者對該場景的效能要求。另外在推理載入模型檔案的場景中,透過預熱的方式,只需要快取叢集從物件儲存上 get一次完整的模型檔案,後面推理節點從快取中讀取,最大程度提升了讀吞吐能力,也節省使用公有云物件儲存的頻寬和呼叫成本。

環節4:混合雲架構資料快速分發的需求

隨著大模型的普及,GPU 算力成為一種稀缺資源,尤其是在國內環境中,不像過去 CPU 資源那樣容易獲取。尤其是需要進行通用模型預訓練的使用者,常常需要處理跨地域和多資料中心算力協作的問題,這就要求儲存資源能夠根據算力的地理位置靈活配置,確保資料能隨算力分佈。

JuiceFS 透過其映象檔案系統功能,實現資料從主站到其他映象站點的自動同步,後設資料同步在確定的時間視窗內完成,同城可以達到毫秒級延遲,跨城市在國內約為 10-30 毫秒,在指定時間視窗內,映象站點保證資料一致性,實際資料複製在後臺非同步完成。一旦後設資料同步完成,資料即可在各站點訪問;如果實際資料尚未完全同步到映象站點,JuiceFS 客戶端會自動將請求路由到主站點的物件儲存以獲取資料,雖有效能損失,但資料訪問不會中斷,並對使用者透明。

在這種混合雲架構中,為了更好地控制成本,建議在一個地點部署物件儲存,各站點本地部署足以容納訓練集資料的快取叢集,以實現近100% 的快取命中率。資料預先從共用物件儲存預熱到各站點的快取叢集中,這種部署方式是AI場景中使用者最廣泛採用的。

下圖是一個使用者的一主三映象站點的混合雲算力叢集架構。每個站點都透過快取組實現了資料訪問的加速。

值得一提的是,在最新企業版中,JuiceFS 新增了映象叢集的可寫功能。在多站點訓練場景中,當各個映象站點的訓練叢集分別儲存 checkpoint 檔案並寫入到各自站點的 JuiceFS 檔案系統時,將自動寫到主站點檔案系統的後設資料和物件儲存,並進行其他映象站點的分發,這裡後設資料的複製是同步的,保證在映象叢集寫資料的版本一致性。這一功能簡化並加快後續訓練任務的 checkpoint 恢復以及推理服務模型載入的流程,更好的保證在映象叢集訓練和推理的資料一致性。映象站點的寫入效能此時會受到站點間專線網路狀況的影響

03 小結

針對大資料量和複雜資料流轉所帶來的挑戰,常見的解決方案包括建立統一的資料管理系統,即資料湖。這種策略透過採用統一的名稱空間來在後端儲存不同型別的資料,同時在前端提供多樣化的資料訪問方式,從而有效降低了複雜資料管道中資料流轉的難度和成本。

在檔案儲存系統的選擇上,使用者可以從基礎架構、效能、相容性、以及易用性等多方面考慮,以選定最適合的產品。JuiceFS 利用廉價的物件儲存作為後端資料持久化手段,並結合多級快取加速架構,為使用者提供提供高成本效益的儲存方案。

此外,文章還探討了在下述關鍵環節, JuiceFS 實踐與最佳化策略。

  • 資料集載入:資料集大檔案和海量小檔案的隨機讀,需要需要檔案系統具備高 IOPS 和低 I/O 延遲,以及高效能的後設資料訪問;
  • 訓練環節 checkpoint 儲存:checkpoint 大檔案的順序寫,需要檔案系統提供高效能的寫吞吐能力;
  • 訓練和推理環節模型載入:safetensors 格式 checkpoint/ 模型大檔案的隨機讀,需要需要檔案系統具備高 IOPS 和低 I/O 延遲;其他格式 checkpoint 模型大檔案的順序讀,需要檔案系統提供高效能的讀吞吐能力;大規模推理節點併發載入模型的頻寬瓶頸問題;
  • 混合雲架構資料分發:可保證資料一致性的映象讀寫功能,滿足多區域協同訓練中的資料集分發、checkpoint 儲存,多區域推理服務的模型部署。

希望這篇內容能夠對你有一些幫助,如果有其他疑問歡迎加入 JuiceFS 社群與大家共同交流。

相關文章