基於文件存儲UFS的Pytorch訓練IO優化實踐

咱們在協助某AI客戶排查一個UFS文件存儲的性能case時發現，其使用的Pytorch訓練IO性能和硬件的IO能力有很大的差距（後面內容有具體性能對比數據）。

讓咱們感到困惑的是： UFS文件存儲，咱們使用fio自測能夠達到單實例最低10Gbps帶寬、IOPS也可達到2w以上。該AI客戶在高IOPS要求的AI單機小模型訓練場景下，或者以前使用MXNet、TensorFlow框架時，IO都能跑到UFS理論性能，甚至在大型分佈式訓練場景中，UFS也能夠徹底勝任。

因而咱們開啓了和客戶的一次深度聯合排查。

初步嘗試優化

1、調整參數：

基於上述狀況，首先考慮是否是使用Pytorch的姿式不對？參考網上提到經驗，客戶調整batch_size、Dataloader等參數。

Batch_size

默認batch_size爲256，根據內存和顯存配置嘗試更改batch_size大小，讓一次讀取數據更多，發現實際對效率沒有提高。經過分析是因爲batch_size設置與數據讀取邏輯沒有直接關係，IO始終會保留單隊列與後端交互，不會下降網絡交互上的總體延時（由於用的是UFS文件存儲，後面會講到爲何用）。

Pytorch Dataloader

Pytorch框架dataloader的worker負責數據的讀取和加載、分配。經過batch_sampler將batch數據分配給對應的worker，由worker從磁盤讀取數據並加載數據到內存，dataloader從內存中讀取相應batch作迭代訓練。這裏嘗試調整了worker_num參數爲CPU核數或倍數，發現提高有限，反而內存和CPU的開銷提高了很多，總體加劇了訓練設備的負擔，經過 worker加載數據時的網絡開銷並不會下降，與本地SSD盤差距依然存在。

這個也不難理解，後面用strace排查的時候，看到CPU更多的時候在等待。

因此：從目前信息來看，調整Pytorch框架參數對性能幾乎沒有影響。

2、嘗試不一樣存儲產品

在客戶調整參數的同時，咱們也使用了三種存儲作驗證，來看這裏是否存在性能差別、差別到底有多大。在三種存儲產品上放上一樣的數據集：

1. 單張平均大小20KB的小圖片，總量2w張。
2. 以目錄樹方式存到三種存儲下的相同路徑，使用Pytorch經常使用的標準讀圖接口CV2和PIL

測試結果，以下圖：

注：SSHFS基於X86物理機（32核/64G/480G SSD*6 raid10）搭建，網絡25Gbps

結論：經過對存儲性能實測， UFS文件存儲較本地盤、單機SSHFS性能差距較大。

爲何會選用這兩種存儲（SSHFS和本地SSD）作UFS性能對比？

當前主流存儲產品的選型上分爲兩類：自建SSHFS/NFS或採用第三方NAS服務（相似UFS產品），個別場景中也會將須要的數據下載到本地SSD盤作訓練。傳統SSD本地盤擁有極低的IO延時，一個IO請求處理基本會在us級別完成，針對越小的文件，IO性能越明顯。受限於單臺物理機配置，沒法擴容，數據基本 「即用即棄」。而數據是否安全也只能依賴磁盤的穩定性，一旦發生故障，數據恢復難度大。可是鑑於本地盤的優點，通常也會用做一些較小模型的訓練，單次訓練任務在較短期便可完成，即便硬件故障或者數據丟失致使訓練中斷，對業務影響一般較小。

用戶一般會使用SSD物理機自建SSHFS/NFS共享文件存儲，數據IO會經過以太網絡，較本地盤網絡上的開銷從us級到ms級，但基本能夠知足大部分業務需求。但用戶須要在平常使用中同時維護硬件和軟件的穩定性，而且單臺物理機有存儲上限，若是部署多節點或分佈式文件系統也會致使更大運維精力投入。

咱們把前面結論放到一塊兒看：

1. 隱形結論：Tensorflow、Mxnet框架無問題。
2. 調整Pytorch框架參數對性能幾乎沒有影響。

三、Pytorch+UFS的場景下， UFS文件存儲較本地SSD盤、單機SSHFS性能差距大。

結合以上幾點信息並與用戶確認後的明確結論：

UFS結合非Pytorch框架使用沒有性能瓶頸， Pytorch框架下用本地SSD盤沒有性能瓶頸，用SSHFS性能可接受。那緣由就很明顯了，就是Pytorch+UFS文件存儲這個組合存在IO性能問題。

深刻排查優化

看到這裏，你們可能會有個疑問：是否是不用UFS，用本地盤就解決了？

答案是不行，緣由是訓練所需的數據總量很大，很容易超過了單機的物理介質容量，另外也出於數據安全考慮，存放單機有丟失風險，而UFS是三副本的分佈式存儲系統，而且UFS能夠提供更彈性的IO性能。

根據以上的信息快速排查3個結論，基本上能夠判斷出： Pytorch在讀UFS數據過程當中，文件讀取邏輯或者UFS存儲IO耗時致使。因而咱們經過strace觀察Pytorch讀取數據總體流程：

經過strace發現，CV2方式讀取UFS裏的文件（NFSV4協議）有不少次SEEK動做，即使是單個小文件的讀取也會「分片」讀取，從而致使了屢次沒必要要的IO讀取動做，而最耗時的則是網絡，從而致使總體耗時成倍增加。這也是符合咱們的猜想。

簡單介紹一下NFS協議特色：

NAS全部的IO都須要通過以太網，通常局域網內延時在1ms之內。以NFS數據交互爲例，經過圖中能夠看出，針對一次完整的小文件IO操做將涉及元數據查詢、數據傳輸等至少5次網絡交互，每次交互都會涉及到client與server集羣的一個TTL，其實這樣的交互邏輯會存在一個問題，當單文件越小、數量越大時則延時問題將越明顯，IO過程當中有過多的時間消耗在網絡交互，這也是NAS類存儲在小文件場景下面臨的經典問題。

對於UFS的架構而言，爲了達到更高擴展性、更便利的維護性、更高的容災能力，採用接入層、索引層和數據層的分層架構模式，一次IO請求會先通過接入層作負載均衡，client端再訪問後端UFS索引層獲取到具體文件信息，最後訪問數據層獲取實際文件，對於KB級別的小文件，實際在網絡上的耗時比單機版NFS/SSHFS會更高。

從Pytorch框架下兩種讀圖接口來看：CV2讀取文件會「分片」進行，而PIL雖然不會「分片」讀取，可是基於UFS分佈式架構，一次IO會通過接入、索引、數據層，網絡耗時也佔比很高。咱們存儲同事也實際測試過這2種方法的性能差別：經過strace發現，相比OpenCV的方式，PIL的數據讀取邏輯效率相對高一些。

優化方向一： 如何下降與UFS交互頻次，從而下降總體存儲網絡延時

CV2：對單個文件而言，「分片讀取」變「一次讀取」

經過對Pytorch框架接口和模塊的調研，若是使用 OpenCV方式讀取文件能夠用2個方法， cv2.imread和cv2.imdecode。

默認通常會用cv2.imread方式，讀取一個文件時會產生9次lseek和11次read，而對於圖片小文件來講屢次lseek和read是沒有必要的。cv2.imdecode能夠解決這個問題，它經過一次性將數據加載進內存，後續的圖片操做須要的IO轉化爲內存訪問便可。

二者的在系統調用上的對好比下圖：

咱們經過使用cv2.imdecode方式替換客戶默認使用的cv2.imread方式，單個文件的總操做耗時從12ms降低到6ms。可是內存沒法cache住過大的數據集，不具有任意規模數據集下的訓練，可是總體讀取性能仍是提高明顯。使用cv2版本的benchmark對一個小數據集進行加載測試後的各場景耗時以下(延遲的非線性降低是由於其中包含GPU計算時間):

PIL：優化dataloader元數據性能，緩存文件句柄

經過PIL方式讀取單張圖片的方式，Pytorch處理的平均延遲爲7ms(不含IO時間)，單張圖片讀取(含IO和元數據耗時)平均延遲爲5-6ms，此性能水平還有優化空間。

因爲訓練過程會進行不少個epoch的迭代，而每次迭代都會進行數據的讀取，這部分操做從屢次訓練任務上來看是重複的，若是在訓練時由本地內存作一些緩存策略，對性能應該有提高。但直接緩存數據在集羣規模上升以後確定是不現實的，咱們初步只緩存各個訓練文件的句柄信息，以下降元數據訪問開銷。

咱們修改了Pytorch的dataloader實現，經過本地內存cache住訓練須要使用的文件句柄，能夠避免每次都嘗試作open操做。測試後發現1w張圖片經過100次迭代訓練後發現，單次迭代的耗時已經基本和本地SSD持平。可是當數據集過大，內存一樣沒法cache住全部元數據，因此使用場景相對有限，依然不具有在大規模數據集下的訓練伸縮性。

UFS server端元數據預加載

以上client端的優化效果比較明顯，可是客戶業務側須要更改少許訓練代碼，最主要是client端沒法知足較大數據量的緩存，應用場景有限，咱們繼續從server端優化，儘可能下降整個鏈路上的交互頻次。

正常IO請求經過負載均衡到達索引層時，會先通過索引接入server，而後到索引數據server。考慮到訓練場景具備目錄訪問的空間局部性，咱們決定加強元數據預取的功能。經過客戶請求的文件，引入該文件及相應目錄下全部文件的元數據，並預取到索引接入server，後續的請求將命中緩存，從而減小與索引數據server的交互，在IO請求到達索引層的第一步便可獲取到對應元數據，從而下降從索引數據server進行查詢的開銷。

通過此次優化以後，元數據操做的延遲較最初可以降低一倍以上，在客戶端不作更改的狀況下，讀取小文件性能已達到本地SSD盤的50%。看來單單優化server端仍是沒法知足預期，經過執行Pytorch的benchmark程序，咱們獲得UFS和本地SSD盤在整個數據讀取耗時。

此時很容易想到一個問題：非Pytorch框架在使用UFS作訓練集存儲時，爲何使用中沒有遇到IO性能瓶頸？

經過調研其餘框架的邏輯發現：不管是MXNet的rec文件，Caffe的LMDB，仍是TensorFlow的npy文件，都是在訓練前將大量圖片小文件轉化爲特定的數據集格式，因此使用UFS在存儲網絡交互更少，相對Pytorch直接讀取目錄小文件的方式，避免了大部分網絡上的耗時。這個區別在優化時給了咱們很大的啓示，將目錄樹級別小文件轉化成一個特定的數據集存儲，在讀取數據作訓練時將IO發揮出最大性能優點。

優化方向二：目錄級內的小文件轉換爲數據集，最大程度降到IO網絡耗時

基於其餘訓練框架數據集的共性功能，咱們UFS存儲團隊趕忙開工，幾天開發了針對Pytorch框架下的數據集轉換工具，將小文件數據集轉化爲UFS大文件數據集並對各個小文件信息創建索引記錄到index文件，經過index文件中索引偏移量可隨機讀取文件，而整個index文件在訓練任務啓動時一次性加載到本地內存，這樣就將大量小文件場景下的頻繁訪問元數據的開銷徹底去除了，只剩下數據IO的開銷。該工具後續也可直接應用於其餘AI類客戶的訓練業務。

工具的使用很簡單，只涉及到兩步：

• 使用UFS自研工具將Pytorch數據集以目錄形式存儲的小文件轉化爲一個大文件存儲到UFS上，生成date.ufs和index.ufs。
• 使用我方提供Folder類替換pytorch原有代碼中的torchvision.datasets.ImageFolder數據加載模塊（即替換數據集讀取方法），從而使用UFS上的大文件進行文件的隨機讀取。只需更改3行代碼便可。

20行：新增from my_dataloader import *

205行：train_dataset = datasets.ImageFolder改成train_dataset = MyImageFolder

224行：datasets.ImageFolder改成MyImageFolder

經過github上Pytorch測試demo對imagenet數據集進行五、十、20小時模擬訓練，分別讀取不一樣存儲中的數據，具體看下IO對總體訓練速度的影響。（數據單位：完成的epoch的個數）

測試條件：

GPU服務器：P404物理機，48核256G，數據盤800G6 SATA SSD RAID10

SSHFS：X86物理機32核/64G，數據盤480G*6 SATA SSD RAID10

Demo：https://github.com/pytorch/examples/tree/master/imagenet

數據集：總大小148GB、圖片文件數量120w以上

經過實際結果能夠看出： UFS數據集方式效率已經達到甚至超過本地SSD磁盤的效果。而UFS數據集轉化方式，客戶端內存中只有少許目錄結構元數據緩存，在100TB數據的體量下，元數據小於10MB，能夠知足任意數據規模，對於客戶業務上的硬件使用無影響。

UFS產品

針對Pytorch小文件訓練場景，UFS經過屢次優化，吞吐性能已獲得極大提高，而且在後續產品規劃中，咱們也會結合現有RDMA網絡、SPDK等存儲相關技術進行持續優化。詳細請訪問：https://docs.ucloud.cn/storage_cdn/ufs/overview

本文做者：UCloud 解決方案架構師 馬傑

歡迎各位與咱們交流有關雲計算的一切~~~