如何藉助 JuiceFS 為 AI 模型訓練提速 7 倍
背景
海量且優質的資料集是一個好的 AI 模型的基石之一,如何儲存、管理這些資料集,以及在模型訓練時提升 I/O 效率一直都是 AI 平臺工程師和演算法科學家特別關注的事情。不論是單機訓練還是分散式訓練,I/O 的效能都會顯著影響整體 pipeline 的效率,甚至是最終的模型質量。
我們也逐漸看到容器化成為 AI 訓練的趨勢,利用容器可以快速彈性伸縮的特點,結合公有云的資源池,能夠最大化資源利用率,為企業大大節約成本。因此也就誕生了類似 Kubeflow [1] 和 Volcano [2] 這樣的開源元件,幫助使用者在 Kubernetes 上管理 AI 任務。Kubernetes 自 1.15 開始新增了 Scheduling Framework [3] ,社群也基於這個新的排程框架優化了很多針對 AI 訓練場景的問題。前面提到的訓練資料管理問題在 Kubernetes 上依然存在,甚至放大了這個需求,因為計算不再是在固定的幾臺機器上進行,資料需要智慧地跟隨計算「流動」(或者反過來)。
最後,不管是演算法科學家日常實驗,還是正式訓練模型,POSIX 介面依然是一個很強烈的需求,雖然主流的框架或者演算法庫基本都支援物件儲存介面但 POSIX 仍然是「第一公民」。一些作業系統的高階特性(如 page cache)也是隻有 POSIX 接口才具備的。
AI 平臺整體架構
上面是一個常見的 AI 平臺架構圖。儲存系統目前使用比較多的就是物件儲存和 HDFS,這裡之所以還會用到 HDFS 有多種原因,比如平臺部署在機房沒有物件儲存,訓練資料集預處理是在大資料平臺等。計算資源混合了 CPU 例項和 GPU 例項,和大資料平臺不一樣的地方在於,AI 平臺的資源天生就是異構的,因此怎麼合理高效利用這些異構資源一直是個業界難題。排程器前面已經介紹到,Kubernetes 是目前主流的元件,結合各種 Job Operator、Volcano、排程外掛可以最大程度上發揮 Kubernetes 的能力。Pipeline 是很重要的一個部分,AI 任務並不只是由模型訓練這一個步驟組成,還包括資料預處理、特徵工程、模型驗證、模型評估、模型上線等多個環節,因此 Pipeline 管理也是非常重要的。最後就是演算法科學家接觸最多的深度學習框架,這些框架目前都有自己的使用群體,很多模型優化會基於某種框架進行(比如 TensorFlow 的 XLA),但也有和框架無關的(比如 TVM [4] )。
本文的關注點在於最底層的儲存層,在保持上層元件不變的情況下,如何優化儲存層的 I/O 效率。這部分包括但不限於資料快取、預讀、併發讀、排程優化等策略,JuiceFS 便是這樣一個儲存層的增強元件,能夠大幅提升 I/O 效率,下面會詳細介紹。
JuiceFS 簡介
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):
得益於元資料快取和資料快取,可以看到不管是在哪種場景下, 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 場景的儲存基石。
引用連結
[1]
Kubeflow: http://www.kubeflow.org
[2]
Volcano: http://volcano.sh
[3]
Scheduling Framework: http://kubernetes.io/docs/concepts/scheduling-eviction/scheduling-framework
[4]
TVM: http://tvm.apache.org
[5]
DLPerf: http://github.com/Oneflow-Inc/DLPerf/tree/master/PyTorch/resnet50v1.5
[6]
warmup
子命令: http://github.com/juicedata/juicefs/blob/main/docs/zh_cn/command_reference.md#juicefs-warmup
[7]
一致性: http://github.com/juicedata/juicefs/blob/main/docs/zh_cn/cache_management.md#%E4%B8%80%E8%87%B4%E6%80%A7
:point_down: 掃碼進群 :point_down:
- JuiceFS 資料讀寫流程詳解
- 如何在 Kubernetes 叢集中玩轉 Fluid JuiceFS
- 如何在 Kubernetes 叢集中玩轉 Fluid JuiceFS
- 如何利用 JuiceFS 的效能工具做檔案系統分析和調優
- 知乎利用 JuiceFS 給 Flink 容器啟動加速實踐
- JuiceFS 在大搜車資料平臺的實踐
- 環球易購資料平臺如何做到既提速又省錢?
- JuiceFS CSI Driver 的最佳實踐
- JuiceFS CSI Driver 的最佳實踐
- 百億級小檔案儲存,JuiceFS 在自動駕駛行業的最佳實踐
- JuiceFS 如何幫助趣頭條超大規模 HDFS 降負載
- ClickHouse 存算分離架構探索
- Elasticsearch 儲存成本省 60%,稿定科技乾貨分享
- Elasticsearch 儲存成本省 60%,稿定科技乾貨分享
- 常見開源分散式檔案系統架構對比
- 如何藉助 JuiceFS 為 AI 模型訓練提速 7 倍
- 如何藉助 JuiceFS 為 AI 模型訓練提速 7 倍
- JuiceFS 使用 TiKV 資料庫上手指南
- JuiceFS CSI Driver v0.10 全新架構解讀
- 基於 JuiceFS 搭建 Milvus 分散式叢集