iGear 用了這個小魔法,模型訓練速度提升 300%

JuiceFS發表於2022-01-28

一個高精度AI模型離不開大量的優質資料集,這些資料集往往由標註結果檔案和海量的圖片組成。在資料量比較大的情況下,模型訓練週期也會相應加長。那麼有什麼加快訓練速度的好方法呢?

壕氣的老闆第一時間想到的通常是提升算力,增加資源。

如果足夠有錢的話,基本不需要再繼續看其他解決方案了。

但大多數情況下,面對昂貴的算力資源,我們不可能無限增加的。那在花了大價錢買到了有限資源的情況下,我們還可以通過什麼方式加快模型訓練,提高資源利用率呢?

本文將為大家介紹的就是 iGear 高效能快取加速方案,我們先看一張簡單的AI工作全流程服務示意圖。

採集的資料通過 iGear 資料中臺的篩選預處理、iGear 標註平臺的標註後會形成優質訓練資料集,這些資料集會流轉到 iGear 訓練平臺來訓練演算法模型。
iGear 訓練平臺基於 Kubernetes 叢集完成異構計算資源排程。在這個架構體系中計算和儲存是分離的,資料集會放在遠端的物件儲存叢集中,當執行模型訓練任務時就需要訪問遠端儲存來獲取資料集,帶來較高的網路 I/O 開銷,也會造成資料集管理不便的問題。

本文提出的 iGear 高效能快取加速方案要做的就是:

1.如何降低 I/O 開銷,提高訓練效率和 GPU 利用率?

2.如何管理資料集,提高使用者易用性和便利性?

快取方案概述

前文已經提及,資料集是儲存在遠端的物件儲存叢集中的。
為了提高資料集的易用性,常用方案是用 fuse 方式掛載給到訓練任務,方便使用者以普通檔案目錄的方式識別並使用資料集。
這種方式雖然滿足了易用性,但高 I/O 的開銷使得使用者必須手動或者用指令碼的方式將資料集提前同步到計算節點,這會增加使用者在訓練過程中的心智負擔。
針對上述問題,我們對訓練資料集做了優化,當使用者開始準備訓練時,通過 JuiceFS 實現資料集快取引擎,為使用者提供資料集快取和預熱功能,這既可以降低對遠端物件儲存的訪問,又可以減少使用者操作。充分利用計算叢集本地儲存來快取資料集,通過兩級快取(訓練節點的系統快取+訓練節點的磁碟快取),加速模型訓練速度,一定程度上也能提高 GPU 的利用率。

JuiceFS 是一款開源的面向雲環境設計的高效能共享檔案系統,在資料組織管理和訪問效能上進行了大量針對性的優化,社群版本也有非常好的文件支撐,這裡就不贅述了。

快取方案測試

測試方案

我們之前使用的方案是使用 fuse 的方式直接掛載物件儲存系統,將 S3 的 bucket 通過掛載的方式掛載到本地,提供訪問遠端物件儲存資料的能力。優化後的高效能快取加速方案,後端也是基於物件儲存,只是在此基礎之上提供了快取、預熱等功能,優化了儲存的效能。
基於此我們做了以下兩組對比實驗,兩組實驗都是基於同一套物件儲存,且其它條件保持一致。

  1. 開啟或關閉高效能快取加速的效能對比

  1. 使用高效能快取加速方案和使用 fuse 掛載的效能對比

測試方式

在伺服器物理機環境下,我們使用 PyTorch /examples 倉庫中提供的 ResNet50 v1.5 進行模型訓練,對其進行單機單卡、單機多卡的結果復現,同時對比執行過程花費的時長。

測試環境

硬體:Tesla V100-SXM2-32GB
驅動:Nvidia 450.80.02
作業系統:Ubuntu 20.04.1 LTS
測試工具:PyTorch ResNet50 v1.5提供的指令碼
資料集:ImageNet

實驗一:開啟高效能快取加速 VS 關閉高效能快取加速

ResNet50 v1.5 batch_size = 128, worker=8

從上圖可以看出,在沒有快取的情況下,訓練任務處理的圖片數並沒有隨著算力的增加而明顯變化,說明已經到了 I/O 瓶頸階段。而使用快取後,隨著算力的增加,處理的圖片數也相應增加。
這證明使用快取加速後,大大降低了 I/O 開銷,同等算力情況下,訓練速度也有了很大提升,其中單機 1 卡的訓練速度提升了 22.3%,單機 4 卡的速度提升了 218%,單機8卡的速度提升了 **230% **。

從模型訓練時間角度來看,訓練時間從未使用快取加速的 1381 min 降低到 565 min,只用了原來 1/3 的時間就完成了模型訓練。

第一組實驗對比了資料集快取前後的效能差異,驗證了使用高效能快取方案來加速 iGear 訓練任務的必要性。

實驗二:高效能快取加速 VS fuse掛載

當前較為通用的方案是採用 fuse 的方式來掛載遠端的物件儲存到本地,提供使用者對資料集的訪問請求。為了對比目前普通 fuse 掛載的方案和高效能快取加速的優化方案,我們設計了第二組實驗:

ResNet50 v1.5 batch_size = 128, worker=8

從模型訓練時間角度來看,fuse 掛載方案訓練時間為 1448 min,高效能快取加速優化方案可將訓練時間減少到 565 min,將近 fuse 掛載方案的 1/3

因此,相比於傳統直接使用物件儲存,我們的高效能儲存在訓練速度上和訓練時間上都有大幅度的提升。

第二組實驗對比了不同方案下的模型訓練時間,驗證了使用高效能快取方案來加速 iGear 訓練任務的重要性。

結論

面對昂貴且有限的算力資源,我們可以通過高效能快取加速方案去大幅加速 iGear 平臺上訓練任務的速度,極大地縮短模型訓練時間,同時提高了 GPU 資源的利用率。如果算力進一步提升,收益也不僅僅止步於當前的測試環境。

如有幫助的話歡迎關注我們專案 Juicedata/JuiceFS 喲! (0ᴗ0✿)

相關文章