作者：星辰算力平臺

1. 背景

隨著大數據、人工智能技術的蓬勃發展，人類對于算力資源的需求也迎來大幅度的增長。在騰訊內部，星辰算力平臺以降本增效為目標，整合了公司的GPU訓練卡資源，為算法工程師們提供統一的底層GPU算力服務。借助于虛擬化、算力挖掘等技術，平臺服務公司內各BG的AI訓練場景，GPU利用率業界領先。同時，通過云原生任務化的方式，對接了內部各大業務，促進了AI技術研究效率的提升和創新研究。

當下，由于AI訓練時的高性能計算設備（如NVIDIA GPU）成本高昂，如果任務在訓練過程中不能保證數據IO的速度，將會導致計算設備低載甚至空載，這無疑在時間和資源上都是一種極大的浪費。

在星辰算力平臺內部，用戶的訓練數據大多存放在平臺提供的CephFS中，訓練時將對應的CephFS目錄掛載至容器內部，從而使用戶在訓練時能夠像使用本地文件系統一樣使用CephFS。但在平臺運營過程中我們發現，在訓練數據集文件數較多時，訓練任務使用CephFS會使訓練速度變得異常緩慢。基于這個普遍存在的問題，本文剖析其產生的原理，然后介紹相應的優化方案。最后，通過延伸思考來發散思維，簡要介紹了不同場景下AI訓練加速的技術。

2. 基本概念

2.1. CephFS IO流程

CephFS IO流程如下圖所示。

CephFS IO路徑

當客戶端進行文件系統調用時（如open、read、readdir等），需要先從元數據服務器（Metadata Server, MDS）中獲取請求文件的元數據信息，元數據信息主要包括文件的Inode號、權限、uid、gid和訪問更改時間等。為了加快元數據的訪問效率，MDS將大部分熱點元數據都緩存在自己的內存中，從而避免低效地通過訪問RADOS（Reliable, Autonomic Distributed Object Store）層來獲取元數據。客戶端在從MDS中獲取元數據后，通過計算的方式（CRUSH算法）得到數據在RADOS中的位置，最后與遠程的存儲設備進行交互。

從這個架構來看，CephFS是一個元數據和用戶數據分離的文件系統。文件的元數據和數據存儲在RADOS中的不同Pool中，客戶端需要先與MDS進行元數據交互，再與RADOS進行數據交互。

2.2. Ceph-FUSE

Ceph-FUSE是CephFS客戶端的一種形式，通過用戶空間文件系統（Filesystem in Userspace, FUSE）的方式來實現CephFS客戶端的功能。FUSE是一個面向類Unix計算機操作系統的軟件接口，它使無特權的用戶能夠無需編輯內核代碼而創建自己的文件系統。目前Linux通過內核模塊對此進行支持。通過這種方式，我們可以編寫用戶態的應用程序，只需要實現Linux定義的一組文件系統接口，即可在用戶態實現一個完整的文件系統。

當用戶需要與CephFS進行交互時，客戶端的整個IO流程如下：

用戶程序通過syscall或glibc庫進行系統調用

進程陷入內核態，文件系統操作請求到達Linux虛擬文件系統（Virtual Filesystem, VFS）

VFS根據請求類型，從Dentry Cache、Inode Cache和Page Cache中分別查找dentry、inode和頁緩存，若緩存命中可直接返回

若緩存不命中，則將請求轉發至FUSE Driver

Ceph-FUSE進程通過libfuse監聽到來自于/dev/fuse的請求，與Ceph集群進行交互并返回結果。

Ceph-FUSE IO路徑

當用戶態程序發起FUSE請求時，Ceph-FUSE在經過處理后會將元數據信息緩存在內存中，提升后續訪問的性能。同時，Linux的Dentry Cache、Inode Cache和Page Cache也會分別緩存該文件的dentry、inode和頁，提升熱點數據的讀取性能。

3. 問題

3.1. 問題源起

星辰算力平臺服務了公司內部各個BG和部門的AI算法工程師，因此平臺上運行的訓練任務場景也各不相同。在運營過程中我們發現，有用戶反映某些任務中CephFS的讀取速度較慢，使整個訓練的時間拉長，其中屬CV類的任務較為明顯。

平臺上CV類的任務數據集，一般都是海量的圖片文件。這類數據集的特點是：

• 文件個數多，小數據集達到十萬級別，大數據集達到百萬、千萬甚至上億級別。

• 單個文件占用空間不大，大多是小文件。

3.2. 理論分析

AI訓練場景與許多復雜的文件操作場景不同，其數據讀寫的邏輯較為簡單。一般來說，用戶會在每個epoch訓練相同的數據，然后訓練多個epoch直至模型達到收斂條件。因此，AI訓練場景下，訓練文件在訓練過程中保持不變，且被讀取的頻率相對固定，同時寫文件的頻率較低。

針對這種特點，由于Ceph-FUSE會對訪問過的元數據進行緩存，同時Linux的Dentry Cache、Inode Cache和Page Cache也會充分緩存讀取過的文件元數據和文件數據。通常來說，在第二個epoch開始時，由于數據集文件在第一個epoch已被訪問過，訓練時的IO速度應當有非常明顯的提升。然而，事與愿違，對于較多數量的文件，我們發現訓練速度沒有明顯提升，且每個epoch的訓練速度都很慢。

為了查出其中的原因，接下來我們復制一個一模一樣的任務，打開Ceph-FUSE日志進行分析。

3.3. 原因排查

3.3.1. Ceph-FUSE日志分析

在訓練任務開始時，打開母機上的Ceph-FUSE日志進行查看。

疑點現象：

在第一個epoch接近末尾時，發現出現了日志trim_caps mds.x max xxx caps xxx。

每次trim_caps執行，清除的dentry個數為5000。

該日志每隔5s會打印一次，往后的訓練過程中會一直持續。

注：CAPS是指capabilities，MDS用CAPS授予客戶端對不同文件進行操作的許可，因此MDS需要實時維護每個客戶端文件操作的CAPS。這就意味著，如果客戶端端持有了某個文件的CAPS并進行了緩存，MDS需要知道每個客戶端緩存了哪些文件。

3.3.2. 提出猜想

根據疑點現象大概能夠提出以下的猜想：

在第一個epoch結束時發生了trim_caps現象，且多次測試結果均是如此，猜測可能是緩存數量到達了某個閾值。

日志每隔5s會打印一次，可能是定時器觸發了trim_caps。

MDS需要維護每個客戶端的CAPS，當客戶端讀取文件數較多時，MDS的cache總會達到oversize的狀態，必定會觸發trim_caps。

3.3.3. 代碼驗證

根據上述猜想，可以在茫茫的Ceph源碼中直奔主題，分別找出MDS和Ceph-FUSE的關鍵代碼。

3.3.3.1. MDS端

根據現象2，在MDS中的tick函數內找到如下代碼：

void?MDSRankDispatcher::tick() {......if?(is_active()?||?is_stopping())?{server->recall_client_state(nullptr,?Server::RecallFlags::ENFORCE_MAX);?//?選中該MDS下持有較多caps數量的客戶端，執行caps回收mdcache->trim();mdcache->trim_client_leases();mdcache->check_memory_usage();?//?當內存使用量過大時，選中該MDS下所有客戶端，執行caps回收(recall_client_state)mdlog->trim();}...... }

從中可以看出，MDS端定時對客戶端的CAPS進行回收，如果回收后內存使用量仍然過高，就對所有客戶端再執行一次CAPS回收。在check_memory_usage函數中會根據cache試用情況決定是否再執行recall_client_state。

void?MDCache::check_memory_usage() {......if?(cache_toofull())?{mds->server->recall_client_state(nullptr);}...... }

進入關鍵函數recall_client_state進行查看。

/***?Call?this?when?the?MDCache?is?oversized,?to?send?requests?to?the?clients*?to?trim?some?caps,?and?consequently?unpin?some?inodes?in?the?MDCache?so*?that?it?can?trim?too.*/ std::pair<bool,?uint64_t>?Server::recall_client_state(MDSGatherBuilder*?gather,?RecallFlags?flags) {......const?bool?enforce_max?=?flags&RecallFlags::ENFORCE_MAX;const?auto?max_caps_per_client?=?g_conf->get_val<uint64_t>("mds_max_caps_per_client");?//?默認為1_Mconst?auto?min_caps_per_client?=?g_conf->get_val<uint64_t>("mds_min_caps_per_client");?//?默認為100const?auto?recall_max_caps?=?g_conf->get_val<uint64_t>("mds_recall_max_caps");?//?默認為5000....../*?trim?caps?of?sessions?with?the?most?caps?first?*/std::multimap<uint64_t,?Session*>?caps_session;auto?f?=?[&caps_session,?enforce_max,?max_caps_per_client](Session*?s)?{auto?num_caps?=?s->caps.size();?//?當前caps總量//?當flags為RecallFlags::ENFORCE_MAX時，只把caps數量超過max_caps_per_client的客戶端找出來，否則找出所有客戶端if?(!enforce_max?||?num_caps?>?max_caps_per_client)?{caps_session.emplace(std::piecewise_construct,?std::forward_as_tuple(num_caps),?std::forward_as_tuple(s));}};mds->sessionmap.get_client_sessions(std::move(f));......for?(const?auto?p?:?boost::adaptors::reverse(caps_session))?{......//?計算每個客戶端的最大caps數量uint64_t?newlim;if?(num_caps?<?recall_max_caps?||?(num_caps-recall_max_caps)?<?min_caps_per_client)?{newlim?=?min_caps_per_client;}?else?{newlim?=?num_caps-recall_max_caps;}if?(num_caps?>?newlim)?{/*?now?limit?the?number?of?caps?we?recall?at?a?time?to?prevent?overloading?ourselves?*/uint64_t?recall?=?std::min<uint64_t>(recall_max_caps,?num_caps-newlim);?//?這里可以看出，每次最多回收mds_recall_max_caps個newlim?=?num_caps-recall;......auto?m?=?new?MClientSession(CEPH_SESSION_RECALL_STATE);?//?新建一個類型為CEPH_SESSION_RECALL_STATE的請求m->head.max_caps?=?newlim;?//?設置客戶端的最大caps數量mds->send_message_client(m,?session);?//?向客戶端發送請求......}......}...... }

從上述代碼基本可以確定CAPS被清除的原因，MDS每隔5s執行了一次recall_client_state。由于mds_max_caps_per_client默認被設置為1_M（也就是1048576），當訓練程序讀取文件個數達到1_M后該客戶端就會被加入caps_session隊列發起CAPS回收請求。由于recall_max_caps默認被設置為5000，所以每次CAPS回收的個數為5000。

3.3.3.2. Ceph-FUSE端

首先，根據MDS端發起的類型為CEPH_SESSION_RECALL_STATE的請求，找到客戶端接受請求的代碼。

void?Client::handle_client_session(MClientSession?*m)? {......switch?(m->get_op())?{......case?CEPH_SESSION_RECALL_STATE:trim_caps(session,?m->get_max_caps());?//?max_caps，值為上述的newlimbreak;......}...... }

Ceph-FUSE接收到MDS的請求后，進入trim_caps函數。

void?Client::trim_caps(MetaSession?*s,?uint64_t?max) {mds_rank_t?mds?=?s->mds_num;size_t?caps_size?=?s->caps.size();?//?客戶端caps總量......uint64_t?trimmed?=?0;auto?p?=?s->caps.begin();std::set<Dentry?*>?to_trim;?//?將需要執行caps回收的Dentry放入其中等待回收//?以下內容通過迭代器p將caps清理至max以下，將需要清理的Dentry放入to_trim中while?((caps_size?-?trimmed)?>?max?&&?!p.end())?{......}for?(const?auto?&dn?:?to_trim)?{trim_dentry(dn);?//?執行Ceph-FUSE內的dentry緩存}to_trim.clear();caps_size?=?s->caps.size();if?(caps_size?>?max)_invalidate_kernel_dcache();?//?這是關鍵函數，調用了Linux的remount操作來清理所有的dentries

Ceph-FUSE接收到MDS的請求后，會將CAPS總量清理至max以下（本例中就是清理5000個CAPS）。同時，將這些CAPS對應的dentry緩存全部清除，并調用操作系統命令來清除Dentry Cache、Inode Cache和Page Cache，執行命令為：

static?int?remount_cb(void?*handle) {//?used?for?trimming?kernel?dcache.?when?remounting?a?file?system,?linux?kernel//?trims?all?unused?dentries?in?the?file?systemchar?cmd[1024];CephFuse::Handle?*cfuse?=?(CephFuse::Handle?*)handle;snprintf(cmd,?sizeof(cmd),?"mount?-i?-o?remount?%s",?cfuse->opts.mountpoint);?//?調用remount，清理文件系統的緩存int?r?=?system(cmd);...... }

3.4. 小結

至此，基本真相大白。整體流程如下圖所示：

訓練程序啟動，開始讀取文件。

在第一個epoch訓練后期，Ceph-FUSE擁有的CAPS達到1_M。

MDS定時器觸發，對持有CAPS超過1_M的客戶端執行發起回收CAPS請求，回收個數為5000。

Ceph-FUSE接收到CEPH_SESSION_RECALL_STATE請求，從caps隊列中清除5000個CAPS并將這些CAPS對應的dentry從cache中清除。

Ceph-FUSE調用Linux的remount命令來清除Linux文件系統的cache。

MDS檢查自身內存使用情況，若超過閾值則重復上述回收操作。

訓練程序第二個epoch后，由于文件系統的cache被清除，導致緩存失效。

CAPS回收流程

4. 解決方案

從上述分析來看，最直觀的改進方法就是將MDS端的參數mds_max_caps_per_client增大，可以使得MDS能夠維護更多的CAPS。然而，這是一種治標不治本的方法。接下來提出一種Ceph-FUSE客戶端緩存的方案，避免客戶端CAPS清除導致訓練速度變慢。

4.1. 元數據緩存方案

4.1.1. 元數據緩存

Ceph針對的是通用場景，設計復雜的CAPS機制來保證多客戶端對同一文件讀寫時的一致性。但在我們的場景中，讀寫方式卻較為固定。主要表現為：

訓練過程中讀取的數據集在訓練過程中不會發生改變，且讀取頻率很高。

寫文件的頻率較低，主要是ckpt和log文件，且不會讀。

在這個特殊的場景下，可以部分犧牲一致性來獲取性能上的提升。具體表現為，Ceph-FUSE側可以將以只讀方式打開的文件進行元數據緩存，減少與MDS的交互，同時在trim_caps發生時不去真正刪除這部分元數據對應的緩存。核心改造如下所示：

當Ceph-FUSE接收到open請求時，如果以只讀方式打開，則將其標記為I_CACHED狀態。在該狀態下的文件操作不會請求MDS獲取CAPS，可以直接從本地cache中讀取元數據，大大減少了與MDS的交互。

如果一個文件被只讀打開后，將無法被讀寫打開，這是為了保證寫數據的一致性。

當trim_caps發生時，Ceph-FUSE將CAPS被回收的Inode標記為I_ORPHAN狀態，然后請求MDS刪除這些CAPS。此時，MDS上已不存在這些Inode的緩存但是本地Ceph-FUSE并沒有真正進行CAPS回收，與此同時也不去清除Linux文件系統的cache，充分保證了元數據的緩存。

元數據緩存方案

以上優化建立的前提是：只讀方式打開的文件不會進行修改。在我們的AI訓練場景下，訓練任務完美契合了這個條件。

4.1.2. 緩存淘汰算法

Ceph-FUSE會將元數據緩存在本地，但其緩存淘汰算法是一種帶高低優先級的LRU算法。LRU算法核心思想是如果數據最近被訪問過，那么將來被訪問的幾率也更高，但這種思想不符合AI訓練的場景。在大多任務訓練過程中，訓練數據文件會被均勻地訪問，每一個epoch中被訪問過的文件反而是這個epoch中不會再被讀取的文件。采用LRU算法會使緩存隊列中即將被用到的文件元數據被刪除，如下圖所示。

LRU淘汰方式

下圖模擬了LRU淘汰策略下訓練數據集命中率分布曲線。

LRU淘汰策略下訓練數據集命中率分布

從該圖中可以看出，LRU淘汰策略下緩存隊列長度越接近數據集大小，命中率提升才越明顯。當隊列長度只有數據集大小的一半時，命中率只有15%左右。

在AI訓練的場景下，采用不替換策略（Not Replacement, NR）將是命中率最高的算法。在訓練的第一個epoch時，Ceph-FUSE將元數據放到緩存中。當緩存隊列已滿時，Ceph-FUSE將不替換現有緩存的數據，保持緩存不變。在第二個epoch時，Ceph-FUSE從緩存隊列中讀取文件元數據，若未命中則請求MDS獲取。

NR算法

4.1.3. 優化結果

結合兩點針對Ceph-FUSE的優化改動，我們對示例任務進行了測試，得到如下的性能測試數據。

訓練任務測試結果

從圖中可以看出，經過優化后針對海量小文件訓練場景，訓練速度的提升非常明顯。在第二個epoch后，元數據緩存優化版本的訓練速度提升為原來的3~4倍，且訓練速度較為穩定。相比于之前的版本，經過優化后的Ceph-FUSE能夠充分利用Linux文件系統的cache，且避免了每個epoch與MDS之間的交互。經過優化后的版本訓練速度能與本地SSD較為貼近。

4.2. 文件緩存方案

文件緩存方案實際上是一種在元數據緩存優化的基礎上，利用本地SSD對文件進行緩存的方案。針對文件數量特別多，利用Linux文件系統cache但是內存不充足的情況，該方法會有一定效果。

訓練程序在第一個epoch訓練時，Ceph-FUSE在處理完read請求后將文件寫入本地SSD中。為了避免海量小文件直接寫入本地造成較多的lookup操作，同時也為了避免任務完成后文件緩存難以進行清理的問題，考慮將所有讀取后的文件進行聚合緩存至一個本地Cache大文件中，由Ceph-FUSE來記錄每個文件在本地Cache文件中的偏移。

文件緩存方案的詳細步驟如下所示：

• 文件緩存命中：

• • 從Metadata Cache中找出文件在本地Cache文件中的偏移。

  • 通過pread從本地SSD緩存文件中讀取指定范圍的字節。

• 文件緩存不命中：

• • 按照正常流程，與Ceph集群進行交互，得到讀取的字節流。

  • 寫本地Cache文件，并記錄該文件在其中的偏移。

  • 更新Metadata Cache，將文件元數據和偏移量加入其中。

文件緩存方案

該方案雖然能夠充分利用本地SSD，但也有一些缺點，具體表現為：

由于第一個epoch讀取文件時，Ceph-FUSE會寫本地Cache文件，可能會使得第一個epoch訓練速度變慢。但當epoch數較多時這部分時間犧牲是值得的。

IO路徑變得更長，Ceph-FUSE需要讀本地文件。

4.3. 方案對比

方案適用場景

原版	在訓練過程中需要修改數據集
元數據緩存	在訓練過程中不修改只讀打開的文件
元數據緩存 + 文件緩存	內存緊張，無法充分使用文件系統緩存

5. 延伸方案

上述分析和方案主要針對的是海量小文件的IO密集型計算場景，接下來發散思維，簡要介紹一下多種AI加速的解決方案。

我們將AI訓練任務分為IO密集型、GPU計算密集型和CPU計算密集型三類任務。

延伸思考

5.1. IO密集型任務

IO密集型任務指的是訓練瓶頸在數據IO上的任務。這類任務一般會讀取較多的數據集文件，數據量較大，GPU由于數據IO的瓶頸一直處于饑餓狀態，因此GPU利用率較低。總結以下幾種解決方案：

元數據緩存

元數據緩存方案能夠將讀取過的文件元數據緩存在內存中。在元數據和用戶數據分離的文件系統中，高效的元數據性能對整個系統性能至關重要。在數據集只讀場景下，元數據緩存可以在FUSE側完成，也可以在用戶側完成。該方案一方面能夠大大較少與元數據服務器之間的交互，緩存熱點元數據，同時也能降級元數據服務器的壓力。

文件緩存

文件緩存方案充分利用了本地SSD進行文件緩存。在數據集只讀場景下，文件緩存仍然是可以在FUSE側完成，也可以在用戶側完成。通過緩存文件元數據并聚合小文件進行本地存儲，能使訓練任務的IO方式從網絡IO逐漸演變為本地IO。

聚合數據集文件

聚合數據集文件方案主要指的是lmdb、TFRecord等技術。在這種方案下，文件數目大大減少，可以有效地緩解深度學習場景下數據存取的問題，進而提高集群資源利用率。但文件聚合存儲的方式對場景有一些限制，比如：數據更新修改會相對麻煩；數據集全局shuffle比較困難，只能做部分的shuffle。

GPUDirect Storage

GPUDirect Storage是NVIDIA公司在2019年推出的有關GPU顯存和存儲設備之間直接進行交互的技術。傳統方式下磁盤中的數據需要先加載至內存中，再拷貝到GPU顯存進行訓練。在這項技術下，可以繞過CPU讓GPU直接與存儲設備進行交互，在本地或遠程存儲（NVMe磁盤）與GPU顯存之間建立直接的數據IO路徑。該方案一方面可以避免主存內數據冗余副本的產生，另一方面也緩解了CPU和內存的壓力。

5.2. GPU計算密集型任務

GPU計算密集型任務指的是訓練瓶頸在GPU計算上的任務，通常需要保證數據IO和梯度同步的低延時，使得GPU時刻處于忙碌狀態。簡要介紹以下幾種解決方案：

數據預取

數據預取是最容易實現的方案。在每一個iteration計算過程中，事先對下一個或幾個iteration所需的數據進行預取并預處理，保證下一個iteration開始時特征已處于就緒狀態。

GPUDirect RDMA(Remote direct memory access)

GPUDirect RDMA從Kepler GPU和CUDA 5.0期間被提出，現在已得到較為廣泛的支持。在多機訓練過程中，這項技術能讓多個GPU之間直接進行通信，同樣也是避免了主存內數據冗余副本的產生，減少數據拷貝環節。配合Mellanox RDMA設備，數據可以從GPU顯存經RDMA網卡發送出去，經另一臺設備的RDMA網卡后傳輸至GPU，大大較少了IO路徑。目前Horovod等分布式訓練工具均以提供對GPUDirect RDMA的支持。

5.3. CPU計算密集型任務

CPU計算密集型任務指的是訓練瓶頸在CPU計算上的任務，這類任務通常的計算瓶頸在于數據的預處理。此類任務CPU處于高負載狀態，但GPU利用率和磁盤IO可能并不高。有以下幾種解決方案：

NVIDIA DALI(Data Loading Library)

NVIDIA DALI是一個經過優化的數據加載的開源庫，提供數據從磁盤加載到訓練的完整pipeline。同時該庫中還提供了音頻處理、圖像處理、視頻處理等預處理方法，能夠將在CPU上執行等預處理步驟放到GPU上快速執行，從而加速AI訓練輸入數據的預處理。

特征存儲

特征存儲方式是一種直觀有效的方案，本質是進行CPU-GPU算力分離。對于某些大規模數據集，事先利用CPU算力對原始數據進行預處理，將樣本特征打包后寫入云存儲設備中，然后多個GPU任務均可共享這些樣本特征數據。但這類方法缺點在于當特征選取發生變化時，需要重新進行預處理。

6. 總結與展望

本文從實際訓練場景出發，首先簡單介紹了CephFS相關的基本概念，接著通過現象和源碼分析訓練過程中讀取文件緩存失效的原因，然后給出了相應的解決方案。經過優化后，測試任務的訓練速度能提升至原來的3~4倍。最后，通過延伸思考來發散思維，簡要介紹了不同場景下AI訓練加速的技術。

未來，針對IO密集型任務，利用GPUDirect Storage和Ceph的RADOS API等技術，結合本地SSD的高速緩存，可以在用戶側探索更極致的加速方案。這種方式理論上能夠擁有更快的文件讀取速度，能在用戶側對文件的元數據和數據進行充分緩存，減少用戶態和內核態轉換，是未來可以繼續研究的方向。

歡迎關注騰訊程序員視頻號

超強干貨來襲 云風專訪：近40年碼齡，通宵達旦的技術人生

總結

以上是生活随笔為你收集整理的海量小文件场景下训练加速优化之路的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。