如何藉助 JuiceFS 為 AI 模型訓練提速 7 倍

語言: CN / TW / HK

背景

海量且優質的資料集是一個好的 AI 模型的基石之一,如何儲存、管理這些資料集,以及在模型訓練時提升 I/O 效率一直都是 AI 平臺工程師和演算法科學家特別關注的事情。不論是單機訓練還是分散式訓練,I/O 的效能都會顯著影響整體 pipeline 的效率,甚至是最終的模型質量。

我們也逐漸看到容器化成為 AI 訓練的趨勢,利用容器可以快速彈性伸縮的特點,結合公有云的資源池,能夠最大化資源利用率,為企業大大節約成本。因此也就誕生了類似 Kubeflow 和 Volcano 這樣的開源元件,幫助使用者在 Kubernetes 上管理 AI 任務。Kubernetes 自 1.15 開始新增了 Scheduling Framework,社群也基於這個新的排程框架優化了很多針對 AI 訓練場景的問題。前面提到的訓練資料管理問題在 Kubernetes 上依然存在,甚至放大了這個需求,因為計算不再是在固定的幾臺機器上進行,資料需要智慧地跟隨計算「流動」(或者反過來)。

最後,不管是演算法科學家日常實驗,還是正式訓練模型,POSIX 介面依然是一個很強烈的需求,雖然主流的框架或者演算法庫基本都支援物件儲存介面但 POSIX 仍然是「第一公民」。一些作業系統的高階特性(如 page cache)也是隻有 POSIX 接口才具備的。

AI 平臺整體架構

image.png 上面是一個常見的 AI 平臺架構圖。儲存系統目前使用比較多的就是物件儲存和 HDFS,這裡之所以還會用到 HDFS 有多種原因,比如平臺部署在機房沒有物件儲存,訓練資料集預處理是在大資料平臺等。計算資源混合了 CPU 例項和 GPU 例項,和大資料平臺不一樣的地方在於,AI 平臺的資源天生就是異構的,因此怎麼合理高效利用這些異構資源一直是個業界難題。排程器前面已經介紹到,Kubernetes 是目前主流的元件,結合各種 Job Operator、Volcano、排程外掛可以最大程度上發揮 Kubernetes 的能力。Pipeline 是很重要的一個部分,AI 任務並不只是由模型訓練這一個步驟組成,還包括資料預處理、特徵工程、模型驗證、模型評估、模型上線等多個環節,因此 Pipeline 管理也是非常重要的。最後就是演算法科學家接觸最多的深度學習框架,這些框架目前都有自己的使用群體,很多模型優化會基於某種框架進行(比如 TensorFlow 的 XLA),但也有和框架無關的(比如 TVM[4])。

本文的關注點在於最底層的儲存層,在保持上層元件不變的情況下,如何優化儲存層的 I/O 效率。這部分包括但不限於資料快取、預讀、併發讀、排程優化等策略,JuiceFS 便是這樣一個儲存層的增強元件,能夠大幅提升 I/O 效率,下面會詳細介紹。

JuiceFS 簡介

image.png JuiceFS 是一個面向雲原生環境設計的高效能開源分散式檔案系統,完全相容 POSIX、HDFS、S3 介面,適用於大資料、AI 模型訓練、Kubernetes 共享儲存、海量資料歸檔管理等場景。

當通過 JuiceFS 客戶端讀取資料時,這些資料將會智慧地快取到應用配置的本地快取路徑(可能是記憶體,也可能是磁碟),同時元資料也會快取到客戶端節點本地記憶體中。對於 AI 模型訓練場景來說,第一個 epoch 完成之後後續的計算都可以直接從快取中獲取訓練資料,極大地提升了訓練效率。

JuiceFS 也具有預讀、併發讀取資料的能力,保證每個 mini-batch 的生成效率,提前準備好資料。

此外 JuiceFS 還提供標準的 Kubernetes CSI Driver,應用可以把 JuiceFS 檔案系統作為一個共享的 Persistent Volume(PV)同時掛載到多個容器中。

得益於以上特性的支援,演算法科學家可以很輕鬆地管理訓練資料,就像訪問本地儲存一樣,無需修改框架進行專門的適配,訓練效率也能得到一定的保障。

測試方案

為了驗證使用 JuiceFS 以後模型訓練的效果,我們選取常見的 ResNet50 模型以及 ImageNet 資料集,訓練任務使用了 DLPerf[5] 專案提供的指令碼,對應的深度學習框架是 PyTorch。訓練節點配置了 8 塊 NVIDIA A100 顯示卡。

作為對比,我們將公有云上的物件儲存作為基準線(通過類 S3FS 的方式進行訪問),同時和開源專案 Alluxio 進行比較,分別測試了 1 機 1 卡、1 機 4 卡、1 機 8 卡不同配置下的訓練效率(即每秒處理的樣本數)。

不論是 JuiceFS 還是 Alluxio,訓練資料集都提前預熱到了記憶體中,資料集約佔用 160G 空間。JuiceFS 提供了 warmup 子命令[6]可以很方便地進行資料集的快取預熱,只需指定需要預熱的目錄或者檔案列表即可。

測試方法是每種配置都跑多輪訓練,每輪只跑 1 個 epoch,將每輪的統計結果彙總,在排除一些可能的異常資料以後,計算得出整體的訓練效率。

JuiceFS 配置選項說明

AI 模型訓練場景的 I/O 模式是典型的只讀模式,即只會對資料集產生讀請求,不會修改資料。因此為了最大化 I/O 效率,可以適當調整一些配置選項(如快取相關配置),下面詳細介紹幾個重要的 JuiceFS 配置選項。

元資料快取

在核心中可以快取三種元資料:屬性(attribute)、檔案項(entry)和目錄項(direntry),它們可以通過如下三個選項控制快取時間: --attr-cache value 屬性快取過期時間;單位為秒 (預設: 1) --entry-cache value 檔案項快取過期時間;單位為秒 (預設: 1) --dir-entry-cache value 目錄項快取過期時間;單位為秒 (預設: 1) 預設元資料在核心中只快取 1 秒鐘,可以根據訓練時長適當增大快取時間,如 2 小時(7200 秒)。

當開啟一個檔案時(即 open() 請求),為了保證一致性[7],JuiceFS 預設都會請求元資料引擎以獲取最新的元資訊。由於資料集都是隻讀的,因此可以適當調整處理策略,設定檢查檔案是否更新的間隔時間,如果時間沒有到達設定的值,則不需要訪問元資料引擎,可以大幅提升開啟檔案的效能。相關配置選項是: --open-cache value 開啟的檔案的快取過期時間(0 代表關閉這個特性);單位為秒 (預設: 0)

資料快取

對於已經讀過的檔案,核心會把它的內容自動快取下來,下次再開啟的時候,如果檔案沒有被更新(即 mtime 沒有更新),就可以直接從核心中的快取(page cache)讀獲得最好的效能。因此當第一個 epoch 執行完畢,如果計算節點的記憶體足夠,那大部分資料集可能都已經快取到 page cache 中,這樣之後的 epoch 將可以不需要經過 JuiceFS 讀取資料,效能也能大幅提升。這個特性已經預設在 0.15.2 及以上版本的 JuiceFS 中開啟,不需要做任何配置。

除了核心中的資料快取,JuiceFS 還支援將資料快取到本地檔案系統中,可以是基於硬碟、SSD 或者記憶體的任意本地檔案系統。本地快取可以通過以下選項來調整: --cache-dir value 本地快取目錄路徑;使用冒號隔離多個路徑 (預設: "$HOME/.juicefs/cache" 或 "/var/jfsCache") --cache-size value 快取物件的總大小;單位為 MiB (預設: 1024) --free-space-ratio value 最小剩餘空間比例 (預設: 0.1) --cache-partial-only 僅快取隨機小塊讀 (預設: false) 例如要將資料快取到記憶體中有兩種方式,一種是將--cache-dir設定為 memory,另一種是將其設定為 /dev/shm。這兩種方式的區別是前者在重新掛載 JuiceFS 檔案系統之後快取資料就清空了,而後者還會保留,效能上兩者沒有太大差別。下面是將資料快取到 /dev/shm/jfscache 並且限定最多使用 300GiB 記憶體的示例: --cache-dir /dev/shm/jfscache --cache-size 307200 JuiceFS 也支援將資料快取到多個路徑,預設會採用輪詢的方式寫入快取資料,多個路徑通過冒號分隔,例如: --cache-dir /data1:/data2:/data3

Alluxio 配置選項說明

Alluxio 的所有元件(如 master、worker、FUSE)都是部署在同一個節點,使用的版本是 2.5.0-2。具體配置如下:

| 配置項 | 設定值 | | ------ | ------ | | alluxio.master.journal.type | UFS | | alluxio.user.block.size.bytes.default | 32MB | | alluxio.user.local.reader.chunk.size.bytes| 32MB | | alluxio.user.metadata.cache.enabled| true | | alluxio.user.metadata.cache.expiration.time | 2day | | alluxio.user.streaming.reader.chunk.size.bytes | 32MB | | alluxio.worker.network.reader.buffer.size | 128MB |

此外 Alluxio FUSE 啟動時指定的掛載選項是:kernel_cache,ro,max_read=131072,attr_timeout=7200,entry_timeout=7200,nonempty

測試結果

測試結果包含兩個場景,一種使用了核心的 page cache,另一種沒有使用。前面提到測試方法是每種配置跑多輪訓練,當跑完第一輪以後,後續的測試都有可能直接從 page cache 中讀取資料。因此我們設計了第二種場景,來測試沒有 page cache 時的訓練效率(比如模型訓練的第一個 epoch),這種場景能更真實反映底層儲存系統的實際效能。

對於第一種場景,JuiceFS 不需要額外配置即可有效利用核心的 page cache,但是物件儲存和 Alluxio 的預設配置都不支援這個特性,需要單獨進行設定。

需要特別注意的是,我們在測試物件儲存的過程中曾經嘗試過開啟 S3FS 的本地快取特性,希望達到類似 JuiceFS 和 Alluxio 的快取效果。但是實際測試時發現即使已經全量預熱快取,以及無論用多少塊顯示卡,1 個 epoch 都無法在 1 天內跑完,甚至比沒有快取時更慢。因此以下測試結果中的「物件儲存」未包含開啟本地快取以後的資料。

下圖是兩個場景的測試結果(「w/o PC」表示沒有 page cache):

image.png

得益於元資料快取和資料快取,可以看到不管是在哪種場景下,JuiceFS 相比物件儲存平均都能達到 4 倍以上的效能提升,最多能有接近 7 倍的效能差距。同時由於物件儲存的訪問方式沒有有效利用核心的 page cache,因此它在這兩種場景的效能差距不大。另外在完整的端到端模型訓練測試中,因為物件儲存的訓練效率太低,跑到指定模型精度所需時間過長,在生產環境中基本屬於不可用狀態。

對比 Alluxio,在有 page cache 的第一種場景中與 JuiceFS 差別不大。在沒有 page cache 只有記憶體快取的第二種場景中 JuiceFS 平均提升 20% 左右的效能,特別是在 1 機 8 卡的配置下,差距進一步加大,達到了 43% 左右的效能差異。Alluxio 在 1 機 8 卡配置下的效能相比 1 機 4 卡沒有提升,沒法充分利用多卡的計算能力。

GPU 資源是一種比較昂貴的資源,因此 I/O 效率的差異也能間接體現到計算資源的成本上,越是能高效利用計算資源才越能整體降低 TCO。

總結及展望

本文介紹了在 AI 模型訓練中如何充分利用 JuiceFS 的特性來為訓練提速,相比直接從物件儲存讀取資料集,通過 JuiceFS 可以帶來最多 7 倍的效能提升。在多卡訓練的場景上也能保持一定的線性加速比,為分散式訓練奠定了基礎。

未來 JuiceFS 還會在 AI 場景探索更多方向,例如進一步提升 I/O 效率、海量小檔案儲存、資料與計算的親和性、與 Job Operator 的結合、與 Kubernetes 排程框架或社群排程器的結合等等。歡迎大家積極參與到 JuiceFS 開源社群中來,共同建設雲原生 AI 場景的儲存基石。