基於 Fluid+JindoCache 加速大模型訓練的實踐

發表於2024-02-28
UI大模型

作者：王濤(揚禮)、陳裘凱(求索)、徐之浩(東伝)

背景

時間步入了 2024 年，新的技術趨勢，如大模型/AIGC/多模態等技術，已經開始與實際業務相結合，並開始生產落地。這些新的技術趨勢不僅提高了算力的需求，也給底層基礎設施帶來了更大的挑戰。

在計算方面，以 GPU 和 FPGA 等異構硬體為例，他們透過短週期的迭代和演進來適應不斷變化的需求。阿里集團透過統一排程、統一資源池以及全面彈性等排程手段滿足了複雜的計算需求。

在儲存方面，經典的微服務應用透過雲原生化的方式，兼顧了效能和效率。但對於計算量增量最大的分散式 AI 訓練、大資料等計算密集型應用，data locality 直接影響了計算作業的執行效率與吞吐，網路 I/O 的消耗還間接拉高了頻寬成本，且在可預見的場景中，資料集規模還會以較高的速率保持增長，如何透過合理的資料快取親和性技術加速資料訪問，將是提升計算任務執行效率的同時降成本的關鍵。

大模型訓練/多媒體等場景的資料集以圖片和音訊檔案為主，天然適合將資料託管在 OSS 物件儲存上，也是目前線上大多數計算作業的儲存選型，以訓練場景為例，具有以下讀資料的特徵：1）資料集順序的隨機化處理造成傳統的單機快取策略失效；2) 多個 epoch 會對資料集進行多輪讀取；3) 作業間可能複用同個資料集；

綜上，阿里巴巴集團內部多個 AI 平臺業務面臨的現狀中，天然適合用分散式快取/檔案系統的形式進行 I/O 層面的加速。

面臨的挑戰

  1. 計算儲存分離架構提升了資料訪問與計算水平擴充套件的靈活度，但導致了資料訪問高延時，對於訓練等對資料快取親和性有顯著訴求的場景延遲不友好：業務團隊使用的機器學習任務在訓練過程中要實時頻繁訪問 OSS 上的資料（以樣本資料集與 checkpoint 為主），在 OSS 頻寬受限或者壓力較大時，訪問 OSS 上資料速度比訪問本地檔案速度要慢 1~2 個數量級，且佔據了使用者大量的頻寬成本；
  2. Kubernetes 排程器資料快取無感知，同一資料來源多次執行訪問依舊慢：在現實應用中深度學習任務執行會不斷重複訪問同一資料，包括相同模型不同超參的任務、微調模型相同輸入的任務、以及 AutoML 任務等。這種深度學習任務的重複資料訪問就產生了可以複用的資料快取。然而，由於原生 Kubernetes 排程器無法感知快取，導致應用排程的結果不佳，快取無法重用，效能難以提升；
  3. OSS 成為資料併發訪問的瓶頸點，穩定性挑戰大：大量機器學習任務在同時訓練時都會併發訪問後端 OSS 儲存。這種併發機器學習訓練造成的 IO 壓力比較大，OSS 服務成為了效能單點，一旦 OSS 頻寬出現瓶頸則會影響所有機器學習任務；
  4. 訓練檔案分散，後設資料壓力：機器學習任務的訓練資料檔案通常會分散在不同路徑下，讀取檔案需要耗費大量的時間在 list 操作上。物件儲存的 list 操作效能較差，在進行大規模 list 時對 OSS 後設資料壓力很大，經常出現超時或者 list 失敗的情況。
  5. IO 穩定性對業務執行有直接影響：導致業務表現不穩定，甚至造成任務失敗。基於 FUSE 的儲存客戶端更容易發生這樣的問題，一旦這些問題無法自動修復，則可能中斷叢集訓練任務。時刻保持 IO 的穩定性是保證業務順利執行的關鍵途徑之一。

在現實應用中，透過對於以上典型資料訪問 pattern 的分析，我們發現 IO 效能問題會導致 GPU 等昂貴計算資源不能被充分利用。機器學習自身訓練的特點導致了資料檔案訪問較分散，後設資料壓力較大。如果能夠精細化地快取後設資料和檔案資料，那麼一方面可以提高快取效率和磁碟利用率，另一方面也可以解決檔案查詢操作帶來的後設資料損耗。

面向深度學習任務的高效快取排程加速系統

為了能更好地滿足集團大規模機器學習模型訓練的高效性需求，模型訓練過程中需要對資料訪問取得更好的資料本地化效果。因此，我們希望達到以下目標：

  • 計算能夠充分利用本地化資料訪問

    避免透過網路反覆讀取，儘量減少 I/O 在計算流水線中的耗時，從而加速機器學習模型的訓練速度，並提升叢集的 GPU 使用率。

  • 降低 OSS 負載壓力

    透過應用對於部分資料的本地讀取，減小資料訪問延時和降低對底層 OSS 的頻寬壓力。

  • 充分發揮熱點資料集的快取節點優勢

    在對使用者無感知的前提下，智慧地將任務排程到資料快取節點上，從而使得常用模型訓練程式越來越快。

  • 後設資料快取和資料快取分離

    可單獨對檔案進行後設資料快取，快取策略定製化。

  • 透過 POSIX 介面讀取資料

    這樣無需在模型開發和訓練階段使用不同的資料訪問介面，降低開發機器學習模型程式的成本。

3.1 架構元件介紹

Fluid

Fluid [ 1] 是一個開源可擴充套件的分散式資料編排和加速系統，以 Kubernetes 標準和對使用者透明的方式為 AI 和大資料等資料密集型應用提供資料訪問能力，其目標為構建雲原生環境下資料密集型應用的高效支撐平臺。

Fluid 透過 Kubernetes 服務提供的資料層抽象，可以讓資料像流體一樣在諸如 HDFS、OSS、Ceph 等儲存源和 Kubernetes 上層雲原生應用計算之間靈活高效地移動、複製、驅逐、轉換和管理。而具體資料操作對使用者透明，使用者不必再擔心訪問遠端資料的效率、管理資料來源的便捷性，以及如何幫助 Kuberntes 做出運維排程決策等問題。使用者只需以最自然的 Kubernetes 原生資料卷方式（PV/PVC）直接訪問抽象出來的資料，剩餘任務和底層細節全部交給 Fluid 處理。

圖片

Fluid 支援多種 Runtime，包括 Jindo，Alluxio，JuiceFS 和 GooseFS；其中能力、效能和穩定性比較突出的是 JindoRuntime，有比較多的真實落地場景。JindoRuntime [ 2] 是 Fluid 一種分散式快取 Runtime 的實現，基於 JindoCache 分散式快取加速引擎。

JindoCache

JindoCache（前身為 JindoFSx）是阿里雲資料湖管理提供的雲原生資料湖加速產品，支援資料快取、後設資料快取等加速功能。JindoCache 能夠為不同檔案路徑使用不同的 CacheSet 從而提供不同的讀寫策略，滿足資料湖的不同使用場景對訪問加速的需求。

JindoCache 可以用於如下場景：

  • OLAP（Presto查詢），提高查詢效能，減少查詢時間。
  • DataServing（HBase），顯著降低 P99 延遲，減少 request 費用。
  • 大資料分析（Hive/Spark 報表），減少報表產出時間，最佳化計算叢集成本。
  • 湖倉一體，減少 request 費用，最佳化 catalog 延遲。
  • AI，加速訓練等場景，減少 AI 叢集使用成本，提供更全面的能力支援。

圖片

KubeDL

一套基於 K8S（ASI）的 AI 工作負載編排系統，負責管理分散式 AI 工作負載的生命週期、與一層排程的互動、容錯與故障恢復、資料集、執行時加速等，高效支撐了集團統一資源池中不同平臺的 AI 訓練任務，包括但不限於淘系、阿里媽媽、達摩院等業務域，日均支撐 1w+ 訓練任務的穩定執行；KubeDL 開源版本 [ 3] 。

專案整體架構圖

圖片

3.2 使用基於 JindoCache 的 Fluid 的原因

  1. Fluid 可以將資料集編排在 Kubernetes 叢集中，實現資料和計算的同置，並且提供基於 Persistent Volume Claim 介面，實現 Kubernetes 上應用的無縫對接。同時 JindoRuntime 提供對 OSS 上資料的訪問和快取加速能力，並且可以利用 FUSE 的 POSIX 檔案系統介面實現可以像本地磁碟一樣輕鬆使用 OSS 上的海量檔案，pytorch 等深度學習訓練工具可利用 POSIX 檔案介面讀取訓練資料。
  2. 提供後設資料和資料分散式快取，可單獨進行後設資料快取預熱。
  3. 提供後設資料快取預熱，避免訓練檔案在 OSS 上大量後設資料操作、提供資料預熱機制，避免在訓練時刻拉取資料造成的資料訪問競爭。
  4. 透過 KubeDL 呼叫 Fluid 資料親和性排程能力，使用者無需感知快取存放的節點位置，以及彈性場景中不斷隨時可能遷移的節點環境，將有資料依賴的任務和已快取的節點進行感知排程，實現儘可能的短路 short-circuit 讀，最大化效能優勢；
  5. JindoCache 提供多種分散式快取能力，可以根據業務需要選擇合適的快取策略。在當前場景中我們選擇 Cache-Aside (Lazy Loading) 的讀快取策略：當應用程式需要讀取資料時，它首先檢查快取以確定資料是否可用。如果資料可用（快取命中），則返回快取的資料。如果資料不可用（快取未命中），則會在底層儲存查詢資料，然後用從底層讀取的資料填充快取，並將資料返回給呼叫者。寫快取策略選擇 Write-Through 即寫時落快取策略，應用程式向底層檔案系統寫入的檔案，同時也會被寫入快取系統中，好處是下一次讀取這部分資料的時候就可以直接從快取系統中讀取，大大提升了讀取效率。
  6. Fluid 支援 FUSE 掛載點自愈能力，可以自動檢查並恢復因 OOM 等異常原因導致的 FUSE 掛載點斷裂問題，避免資料訪問異常，保障 AI 平臺線上業務穩定執行。

3.3 落地實踐

在集團場景的實踐中，我們基於 KubeDL 的作業編排能力，結合 Fluid+JindoRuntime 的快取引擎能力，同時充分利用了集團龐大異構計算資源池中閒置的記憶體/高效能磁碟等本地資源，端到端地為 AI 平臺提供了資料 I/O 的加速能力。

  1. 集團龐大的統一異構資源池提供了差異化 SLO 的資源售賣等級，執行著高保障、Spot Instance、潮汐離線、普通離線 等多種不同等級的資源，以及搭配了多種代系的機型、SSD、高效能網路卡等硬體，透過合理搭配 JindoCache 的多級快取介質，我們能充分利用好統一資源池的閒置資源；
  2. 結合 JindoCache 快取叢集的組成特點，我們使用高保障的計算資源執行後設資料服務，利用彈性的離線資源來執行 Cache 快取節點服務（IO  Bound 型別），在充分結合了集團資源池排程特點的同時最小化了使用者成本；
  3. 結合 KubeDL 的分散式訓練任務管理與 Fluid 資料集管理能力，我們針對不同使用者的相同資料來源自動進行資料集的跨作業複用，甚至不同平臺的相同資料來源也可以在統一資源池中自動複用，並且基於作業的引用計數，KubeDL 可以自動回收閒置的資料集以幫助使用者主動節約成本。

3.4 經驗分享

根據實踐，我們總結了以下五個方面的經驗供大家參考。

1. 選擇合適的快取節點

使用 JindoRuntime 可以獲得更好的資料本地效能，在實際生產中我們發現不是所有節點都來做快取效能就比較好。原因是有些節點的磁碟和網路 IO 效能不是很好，這個時候需要我們能夠把快取節點儘量選擇到一些大容量磁碟和網路較好的節點上。Fluid 支援 dataset 的可排程性，換言之，就是快取節點的可排程性，我們透過指定 dataset 的 nodeAffinity 來進行資料集快取節點的排程，從而保證快取節點可高效的提供快取服務。

2. 配置快取容量與路徑

透過 dataset 的 Mounts 和 JindoRuntime 的 tieredstore 可以設定資料的掛載目錄。同時，為避免資料量過多而導致快取量過於龐大，可手動配置 JindoRuntime 的 tieredstore 來約束快取的最大容量與水位線（超過水位線的資料會被自動丟棄），tieredstore 也包含對快取存放路徑的設定與儲存層（SSD/MEM/HDD）的設定，以滿足各種場景的需要。對於多節點的場景，使用dataset 的 replacement 可以支援在同一叢集上部署多個 dataset。

3. 設定快取安全策略

在 Fluid 中建立 Dataset 時，有時候我們需要在 mounts 中配置一些敏感資訊，如 OSS 賬號的 accessKeyId、accessKeySecret 。為了保證安全，Fluid 提供使用 Secret 來配置這些敏感資訊的能力。透過建立 Secret，dataset 以 EncryptOptions 欄位指定 Secret 的 name，實現對敏感資訊的繫結。

4. 資料預載入

對於已經建立完成的 dataset 和 jindoruntime，第一次訪問掛載的資料會經歷一次下載資料目錄下全部檔案的過程，這就產生了一個問題：若資料所在的目錄存在無需使用的其他資料，會造成無意義的空間資源與網路資源浪費。為避免這種問題，Fluid 既支援對資料的預載入，同時也支援後設資料快取。透過建立 dataload 讀取所要預載入資料路徑資訊，可以動態將資料注入。dataload 支援快取後設資料與遮蔽非預載入資料的訪問，這樣就大大降低的資料訪問效率。

5. 啟用 Fluid FUSE 掛載點自愈能力

線上業務執行過程中，FUSE 程式可能因為記憶體資源不足以及其他原因崩潰重啟，造成業務容器內 FUSE 掛載點斷聯，出現資料訪問異常並影響線上業務可用性。透過啟用 Fluid FUSE 掛載點自愈能力，Fluid 自動檢測並修復此類斷聯掛載點，持續保障線上業務穩定執行。

3.5 實踐結果

讀樣本加速

以生產環境中的真實使用者作業為基礎，我們對 JindoCache 的效果進行了一次端到端的驗證。

  • 目標任務： LLAMA 13B 的預訓練任務
  • 實驗環境：
  • 叢集&機型：高效能網路叢集 A800 伺服器，搭載 RDMA 網路卡與 Nvme 高速硬碟；
  • JindoCache 規格：預設值為 24*32Gi Cache Worker，以 Nvme 盤為儲存介質（相對記憶體的價效比更高）。

實驗結論：

LLaMa 13B 預訓練模型

I/O訪問模式GPU UtilSM UtilTFLops（log）TFLops（amperf）
直連100%~60%~135~60（avg 10m）
JindoCache快取100%~80%（↑33%）~160（↑18%）~72（avg 10m）（↑20%）

監控資料-無快取直連

圖片

圖片

圖片

監控資料-開啟快取

圖片

整機的平均 GPU 利用率同樣接近 100%，但是各卡的負載較為均勻，都接近 100%。

圖片

Checkpoint 加速

訓練/離線推理場景

分散式訓練任務在每次重啟任務時都會 load checkpoint 模型檔案以繼續訓練，模型大小從幾百 MB 到幾十 GB 不等；除此之外還有大量的離線推理任務，大量使用了統一資源池中的 Spot Instance 彈性 GPU 資源，每個推理任務都會隨時被搶佔，並在 FailOver 之後重新載入模型做離線推理，因此會有大量 Job 在“生生滅滅”中載入同一個 Checkpoint 檔案。

透過 JindoCache 的分散式快取加速，我們將“多次遠端讀”變成了“單次本地讀”，極大加速了 Job FailOver 速度的同時還為使用者節約了多次反覆讀的頻寬成本，在典型的大模型場景中，7b 引數量搭配 fp16 精度，模型檔案的體積約 20Gb，透過 JindoCache 加速我們將使用者每次載入模型的耗時從 10min 縮短到了約 30s。

圖片

訓練 Spot 場景（寫時落快取）

分散式訓練的 Spot 場景中，同步訓練任務通常會在被搶佔之後 FailOver 重新全域性重啟並續跑，KubeDL 會與一層排程配合，以互動式搶佔的方式通知到訓練任務的 Rank 0 節點做一次 on-demand checkpoint 以儲存最新的訓練進度，並能夠在重啟後 reload 最新的 checkpoint 及時續跑，享受 Spot 彈性低成本的同時最小化訓練中斷的代價。

透過最新版本的 JindoCache 寫時落快取能力，原先重新後從遠端重新被動 load 最新的模型檔案，變成了重啟後即時從本地快取叢集 load 最新的模型檔案，端到端 FailOver 的中止時間從平均 10min 縮短到了平均 2min，節約了 80% 的閒置寶貴算力損失。

總結與展望

綜上，使用基於 JindoRuntime 的 Fluid 在集團大規模機器學習模型訓練中發揮了重要作用。在讀樣本加速場景中，透過使用 JindoCacheEngine，我們大大提升了系統的吞吐，使得 GPU 負載利用更加均衡。同時，JindoRuntime 的抽象層遮蔽了不同 JindoCache 版本之間的差異，從而實現了無縫的升級。在 CheckPoint 加速環節，端到端模型載入速度顯著提升，我們以較低成本實現了效能的大幅提升。

未來，我們將繼續透過 Fluid 進行更多場景的嘗試，以及對現有功能的擴充，例如：

  1. 基於引用計數，自動回收閒置資料集（DataSet），實現多資料集的智慧管理。
  2. 智慧資料預熱，基於任務訪問資料模式的自動預熱，按目錄優先順序預熱/驅逐，並進行並行預熱（按目錄拆解預熱任務)。
  3. 啟用 RDMA 技術來加速叢集內的 worker 傳輸吞吐，從而充分利用叢集的高效能網路基礎設施。

最後，在充分利用 JindoCache 快取加速能力和 Fluid 的多 JindoCache 編排能力的基礎上，我們將對使用方式和上層系統的接入進行最佳化，推動硬體和軟體協同，進一步提升效能並支援新硬體。

相關連結：

[1] Fluid

https://github.com/fluid-cloudnative/fluid

[2] JindoCache

https://github.com/aliyun/alibabacloud-jindodata/blob/master/docs/user/6.x/6.2.0/jindo_fluid/jindo_fluid_overview.md

[3] KubeDL 開源版本

https://github.com/kubedl-io/kubedl

如果你對 Fluid 專案感興趣，歡迎點選此處瞭解更多。

相關文章