我們在協助某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 DataloaderPytorch架構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 | 本地SSD | UFS檔案存儲 |
| CV2 | 319.94張/s | 554.73張/s | 72.41張/s |
| PIL（image_open） | 435.93張/s | 3507.93張/s | 115.78張/s |
注：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架構參數對性能幾乎沒有影響。
3、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互動頻次，進而降低整體存儲網絡延時
1、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計算時間):
| UFS imread方式 | UFS imdecode方式 | |
| 95s | 150s | 100s左右 |
2、PIL：優化dataloader中繼資料性能，緩存檔案句柄
通過PIL方式讀取單張圖檔的方式，Pytorch處理的平均延遲為7ms(不含IO時間)，單張圖檔讀取(含IO和中繼資料耗時)平均延遲為5-6ms，此性能水準還有優化空間。由于訓練過程會進行很多個epoch的疊代，而每次疊代都會進行資料的讀取，這部分操作從多次訓練任務上來看是重複的，如果在訓練時由本地記憶體做一些緩存政策，對性能應該有提升。但直接緩存資料在叢集規模上升之後肯定是不現實的，我們初步隻緩存各個訓練檔案的句柄資訊，以降低中繼資料通路開銷。我們修改了Pytorch的dataloader實作，通過本地記憶體cache住訓練需要使用的檔案句柄，可以避免每次都嘗試做open操作。測試後發現1w張圖檔通過100次疊代訓練後發現，單次疊代的耗時已經基本和本地SSD持平。但是當資料集過大，記憶體同樣無法cache住所有中繼資料，是以使用場景相對有限，依然不具備在大規模資料集下的訓練伸縮性。
| 存儲類型 | Real time（s） | User time（s） | Sys time（s） |
| 25.10 | 475.89 | 3.66 | |
| UFS | 25.90 | 491.62 | 3.73 |
3、UFS server端中繼資料預加載
以上client端的優化效果比較明顯，但是客戶業務側需要更改少量訓練代碼，最主要是client端無法滿足較大資料量的緩存，應用場景有限，我們繼續從server端優化，盡量降低整個鍊路上的互動頻次。正常IO請求通過負載均衡到達索引層時，會先經過索引接入server，然後到索引資料server。考慮到訓練場景具有目錄通路的空間局部性，我們決定增強中繼資料預取的功能。通過客戶請求的檔案，引入該檔案及相應目錄下所有檔案的中繼資料，并預取到索引接入server，後續的請求将命中緩存，進而減少與索引資料server的互動，在IO請求到達索引層的第一步即可擷取到對應中繼資料，進而降低從索引資料server進行查詢的開銷。經過這次優化之後，中繼資料操作的延遲較最初能夠下降一倍以上，在用戶端不做更改的情況下，讀取小檔案性能已達到本地SSD盤的50%。看來單單優化server端還是無法滿足預期，通過執行Pytorch的benchmark程式，我們得到UFS和本地SSD盤在整個資料讀取耗時。
| 1w張耗時（s） | 2w張耗時（s） | |
| 48.86 | 97.06 | |
| 97.98 | 195.82 |
此時很容易想到一個問題：非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資料集進行5、10、20小時模拟訓練，分别讀取不同存儲中的資料，具體看下IO對整體訓練速度的影響。（資料機關：完成的epoch的個數）
| 資料讀取方式 | 第一次（5小時） | 第二次（10小時） | 第三次（20小時） |
| 本地SSD盤 | 5.226 | 10.384 | 20.873 |
| UFS目錄小檔案 | 1.238 | 3.180 | 6.394 |
| UFS資料集 | 5.46 | 10.739 | 21.784 |
| SSHFS目錄小檔案 | 3.721 | 7.398 | 14.797 |
測試條件：GPU伺服器：P40*4實體機，48核256G，資料盤800G*6 SATA SSD RAID10SSHFS：X86實體機32核/64G，資料盤480G*6 SATA SSD RAID10Demo：
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。