如何藉助 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: