導讀 近些年，随着深度學習的崛起， Alluxio 分布式緩存技術逐漸成為業界解決雲上 IO 性能問題的主流方案。不僅如此，Alluxio 還天然具備資料湖所需的統一管理和通路的能力。本文将分享面向大規模深度學習訓練的緩存優化，主要分析如今大規模深度學習訓練的存儲現狀與挑戰，說明緩存資料編排在深度學習訓練中的應用，并介紹大規模緩存系統的資源配置設定與排程。
全文目錄如下：
1. 項目背景和緩存政策
2. SiloD 架構
3. 分布式緩存政策以及副本管理
分享嘉賓｜張虔熙 微軟 進階研發工程師
編輯整理｜任文強 xtransfer
出品社群｜DataFun
01
項目背景和緩存政策
首先來分享一下相關背景。
近年來，AI 訓練應用越來越廣泛。從基礎架構角度來看，無論是大資料還是 AI 訓練叢集中，大多使用存儲與計算分離的架構。比如很多 GPU 的陣列放到一個很大的計算叢集中，另外一個叢集是存儲。也可能是使用的一些雲存儲，像微軟的 Azure 或者是亞馬遜的 S3 等。
這樣的基礎架構的特點是，首先，計算叢集中有很多非常昂貴的 GPU，每台 GPU 往往有一定的本地存儲，比如 SSD 這樣的幾十 TB 的存儲。這樣一個機器組成的陣列中，往往是用高速網絡去連接配接遠端，比如 Coco、 image net、YouTube 8M 之類的非常大規模的訓練資料是以網絡進行連接配接的。
如上圖所示，資料有可能會成為下一個 AI 訓練的瓶頸。我們觀察到資料集越來越大，随着 AI 應用更加廣泛，也在積累更多的訓練資料。同時 GPU 賽道是非常卷的。比如 AMD、TPU 等廠商，花費了大量精力去優化硬體和軟體，使得加速器，類似 GPU、TPU這些硬體越來越快。随着公司内加速器的應用非常廣泛之後，叢集部署也越來越大。這裡的兩個表呈現了關于資料集以及 GPU 速度的一些變化。之前的 K80 到 V100、 P100、 A100，速度是非常迅速的。但是，随着速度越來越快，GPU 變得越來越昂貴。我們的資料，比如 IO 速度能否跟上 GPU 的速度，是一個很大的挑戰。
如上圖所示，在很多大公司的應用中，我們觀察到這樣一個現象：在讀取遠端資料的時候，GPU 是空閑的。因為 GPU 是在等待遠端資料讀取，這也就意味着 IO 成為了一個瓶頸，造成了昂貴的 GPU 被浪費。有很多工作在進行優化來緩解這一瓶頸，緩存就是其中很重要的一個優化方向。這裡介紹兩種方式。
第一種，在很多應用場景中，尤其是以 K8s 加 Docker 這樣的基礎 AI 訓練架構中，用了很多本地磁盤。前文中提到 GPU 機器是有一定的本地存儲的，可以用本地磁盤去做一些緩存，把資料先緩存起來。
啟動了一個 GPU 的 Docker 之後，不是馬上啟動 GPU 的 AI 訓練，而是先去下載下傳資料，把資料從遠端下載下傳到 Docker 内部，也可以是挂載等方式。下載下傳到 Docker 内部之後再開始訓練。這樣盡可能的把後邊的訓練的資料讀取都變成本地的資料讀取。本地 IO 的性能目前來看是足夠支撐 GPU 的訓練的。VLDB 2020 上面，有一篇 paper，CoorDL，是基于 DALI 進行資料緩存。
這一方式也帶來了很多問題。首先，本地的空間是有限的，意味着緩存的資料也是有限的，當資料集越來越大的時候，很難緩存到所有資料。另外，AI 場景與大資料場景有一個很大的差別是，AI 場景中的資料集是比較有限的。不像大資料場景中有很多的表，有各種各樣的業務，每個業務的資料表的内容差距是非常大的。在 AI 場景中，資料集的規模、資料集的數量遠遠小于大資料場景。是以常常會發現，公司中送出的任務很多都是讀取同一個資料。如果每個人下載下傳資料到自己本地，其實是不能共享的，會有非常多份資料被重複存儲到本地機器上。這種方式顯然存在很多問題，也不夠高效。
接下來介紹第二種方式。既然本地的存儲不太好，那麼，是否可以使用像 Alluxio 這樣一個分布式緩存來緩解剛才的問題，分布式緩存有非常大的容量來裝載資料。另外，Alluxio 作為一個分布式緩存，很容易進行共享。資料下載下傳到 Alluxio 中，其他的用戶端，也可以從緩存中讀取這份資料。這樣看來，使用 Alluxio 可以很容易地解決上面提到的問題，為 AI 訓練性能帶來很大的提升。微軟印度研究院在 FAST2020 發表的名為 Quiver 的一篇論文，就提到了這樣的解決思路。但是我們分析發現，這樣一個看似完美的配置設定方案，還是比較靜态的，并不高效。同時，采用什麼樣的 cache 淘汰算法，也是一個很值得讨論的問題。
如上圖所示，是使用 Alluxio 作為 AI 訓練的緩存的一個應用。使用 K8s 做整個叢集任務的排程和對 GPU、CPU、記憶體等資源的管理。當有使用者送出一個任務到 K8s 時，K8s 首先會做一個插件，通知 Alluxio 的 master，讓它去下載下傳這部分資料。也就是先進行一些熱身，把作業可能需要的任務，盡量先緩存一些。當然不一定非得緩存完，因為Alluxio 是有多少資料，就使用多少資料。剩下的，如果還沒有來得及緩存，就從遠端讀取。
另外，Alluxio master 得到這樣的指令之後，就可以讓排程它的 worker 去遠端。可能是雲存儲，也可能是 Hadoop 叢集把資料下載下傳下來。這個時候，K8s 也會把作業排程到 GPU 叢集中。比如上圖中，在這樣一個叢集中，它選擇第一個節點和第三個節點啟動訓練任務。啟動訓練任務之後，需要進行資料的讀取。在現在主流的像 PyTorch、Tensorflow 等架構中，也内置了 Prefetch，也就是會進行資料預讀取。它會讀取已經提前緩存的 Alluxio 中的緩存資料，為訓練資料 IO 提供支援。當然，如果發現有一些資料是沒有讀到的，Alluxio 也可以通過遠端進行讀取。Alluxio 作為一個統一的接口是非常好的。同時它也可以進行資料的跨作業間的共享。
如上圖所示，比如又有一個人送出了同樣資料的另一個作業，消耗的是同一個資料集，這個時候，當送出作業到 K8s 的時候，Alluxio 就知道已經有這部分資料了。如果 Alluxio 想做的更好，甚至是可以知道，資料即将會被排程到哪台機器上。比如這個時候排程到 node 1、node 3 和 node 4 上。node 4 的資料，甚至可以做一些副本進行拷貝。這樣所有的資料，即使是 Alluxio 内部，都不用跨機器讀，都是本地的讀取。是以看起來 Alluxio 對 AI 訓練中的 IO 問題有了很大的緩解和優化。但是如果仔細觀察，就會發現兩個問題。
第一個問題就是緩存的淘汰算法非常低效，因為在 AI 場景中，通路資料的模式跟以往有很大差別。第二個問題是，緩存作為一種資源，與帶寬（即遠端存儲的讀取速度）是一個對立的關系。如果緩存大，那麼從遠端讀取資料的機會就小。如果緩存很小，則很多資料都得從遠端讀取。如何很好地排程配置設定這些資源也是一個需要考慮的問題。
在讨論緩存的淘汰算法之前，先來看一下 AI 訓練中資料通路的過程。在 AI 訓練中，會分為很多個 epoch，不斷疊代地去訓練。每一個訓練 epoch，都會讀取每一條資料，并且僅讀一次。為了防止訓練的過拟合，在每一次 epoch 結束之後，下一個 epoch 的時候，讀取順序會變化，會進行一個 shuffle。也就是每次每個 epoch 都會把所有資料都讀取一次，但是順序卻不一樣。
Alluxio 中預設的 LRU 淘汰算法，顯然不能很好地應用到AI訓練場景中。因為 LRU 是利用緩存的本地性。本地性分為兩方面，首先是時間本地性，也就是現在通路的資料，馬上可能還會即将通路。這一點，在 AI 訓練中并不存在。因為現在通路的資料，在下一輪的時候才會通路，而且下一輪的時候都會通路。沒有一個特殊的機率，一定是比其他資料更容易被通路。另一方面是資料本地性，還有空間本地性。也就是，為什麼 Alluxio 用比較大的 block 緩存資料，是因為某條資料讀取了，可能周圍的資料也會被讀取。比如大資料場景中，OLAP 的應用，經常會進行表的掃描，意味着周圍的資料馬上也會被通路。但是在 AI 訓練場景中是不能應用的。因為每次都會 shuffle，每次讀取的順序都是不一樣的。是以 LRU 這種淘汰算法并不适用于 AI 訓練場景。
不僅是 LRU，像 LFU 等主流的淘汰算法，都存在這樣一個問題。因為整個 AI 訓練對資料的通路是非常均等的。是以，可以采用最簡單的緩存算法，隻要緩存一部分資料就可以，永遠不用動。在一個作業來了以後，永遠都隻緩存一部分資料。永遠都不要淘汰它。不需要任何的淘汰算法。這可能是目前最好的淘汰機制。
如上圖中的例子。上面是 LRU 算法，下面是均等方法。在開始隻能緩存兩條資料。我們把問題簡單一些，它的容量隻有兩條，緩存 D 和 B 這兩條資料，中間就是通路的序列。比如命中第一個通路的是 B，如果是 LRU，B 存在的緩存中命中了。下一條通路的是 C，C 并不在 D 和 B，LRU 的緩存中，是以基于 LRU 政策，會把 D 替換掉，C 保留下來。也就是這個時候緩存是 C 和 B。下一個通路的是 A，A 也不在 C 和 B 中。是以會把B 淘汰掉，換成 C 和 A。下一個就是 D，D 也不在緩存中，是以換成 D 和 A。以此類推，會發現所有後面的通路，都不會再命中緩存。原因是在進行 LRU 緩存的時候，把它替換出來，但其實在一個 epoch 中已經被通路一次，這個 epoch 中就永遠不會再被通路到了。LRU 反倒把它進行緩存了，LRU 不但沒有幫助，反倒是變得更糟糕了。不如使用 uniform，比如下面這種方式。
下面這種 uniform 的方式，永遠在緩存中緩存 D 和 B，永遠不做任何的替換。在這樣情況下，你會發現至少有 50% 的命中率。是以可以看到，緩存的算法不用搞得很複雜，隻要使用 uniform 就可以了，不要使用 LRU、LFU 這類算法。
對于第二個問題，也就是關于緩存和遠端帶寬之間關系的問題。現在所有主流的 AI 架構中都内置了資料預讀，防止 GPU 等待資料。是以當 GPU 做訓練的時候，其實是觸發了 CPU 預取下一輪可能用到的資料。這樣可以充分利用 GPU 的算力。但當遠端存儲的 IO 成為瓶頸的時候，就意味着 GPU 要等待 CPU 了。是以 GPU 會有很多的空閑時間，造成了資源的浪費。希望可以有一個比較好的排程管理方式，緩解 IO 的問題。
緩存和遠端 IO 對整個作業的吞吐是有很大影響的。是以除了 GPU、CPU 和記憶體，緩存和網絡也是需要排程的。在以往大資料的發展過程中，像 Hadoop、yarn、my source、K8s 等，主要都是排程 CPU、記憶體、GPU。對于網絡，尤其對于緩存的控制都不是很好。是以，我們認為，在 AI 場景中，需要很好的排程和配置設定它們，來達到整個叢集的最優。
02
SiloD 架構
在 EuroSys 2023 發表了這樣一篇文章，它是一個統一的架構，來排程計算資源和存儲資源。
整體架構如上圖所示。左下角是叢集中的 CPU 和 GPU 硬體計算資源，以及存儲資源，如 NFS、雲存儲 HDFS 等。在上層有一些 AI 的訓練架構 TensorFlow、PyTorch 等。我們認為需要加入一個統一管理和配置設定計算和存儲資源的插件，也就是我們提出的 SiloD。
如上圖所示，一個作業可以達到什麼樣的吞吐和性能，是由 GPU 和 IO 的最小值決定的。使用多少個遠端 IO，就會使用多少遠端的 networking。可以通過這樣一個公式算出通路速度。作業速度乘以緩存未命中率，也就是（1-c/d）。其中 c 就是緩存的大小，d 就是資料集。這也就意味着資料隻考慮 IO 可能成為瓶頸的時候，大概的吞吐量是等于（b/(1-c/d)），b 就是遠端的帶寬。結合以上三個公式，可以推出右邊的公式，也就是一個作業最終想達到什麼樣的性能，可以這樣通過公式去計算沒有 IO 瓶頸時的性能，和有 IO 瓶頸時的性能，取二者中的最小值。
得到上面的公式之後，把它微分一下，就可以得到緩存的有效性，或者叫做緩存效率。即雖然作業很多，但在配置設定緩存的時候不能一視同仁。每一個作業，基于資料集的不同，速度的不同，緩存配置設定多少是很有講究的。這裡舉一個例子，就以這個公式為例，如果發現一個作業，速度非常快，訓練起來非常快，同時資料集很小，這時候就意味着配置設定更大的緩存，收益會更大。
基于以上觀察，可以使用 SiloD，進行緩存和網絡的配置設定。而且緩存的大小，是針對每個作業的速度，以及資料集整個的大小來進行配置設定的。網絡也是如此。是以整個架構是這樣的：除了主流的像 K8s 等作業排程之外，還有資料管理。在圖左邊，比如緩存的管理，要統計或者監控配置設定整個叢集中緩存的大小，每個作業緩存的大小，以及每個作業使用到的遠端 IO 的大小。底下的作業，和 Alluxio 方式很像，都可以都使用 API 進行資料的訓練。每個 worker 上使用緩存對于本地的 job 進行緩存支援。當然它也可以在一個叢集中跨節點，也可以進行共享。
經過初步測試和實驗，發現這樣一個配置設定方式可以使整個叢集的使用率和吞吐量都得到非常明顯的提升，最高可以達到 8 倍的性能上的提升。可以很明顯的緩解作業等待、GPU 空閑的狀态。
對上述介紹進行一下總結：
第一，在 AI 或者深度學習訓練場景中，傳統的 LRU、LFU 等緩存政策并不适合，不如直接使用 uniform。
第二，緩存和遠端帶寬，是一對夥伴，對整體性能起到了非常大的作用。
第三，像 K8s、yarn 等主流排程架構，可以很容易繼承到 SiloD。
最後，我們在 paper 中做了一些實驗，不同的排程政策，都可以帶來很明顯的吞吐量的提升。
03
分布式緩存政策以及副本管理
我們還做了一些開源的工作。分布式緩存政策以及副本管理這項工作，已經送出給社群，現在處于 PR 階段。Alluxio master 主要做 Meta 的管理和整個 worker 叢集的管理。真正緩存資料的是 worker。上面有很多以 block 為機關的塊兒去緩存資料。存在的一個問題是，現階段的緩存政策都是單個 worker 的，worker 内部的每個資料在進行是否淘汰的計算時，隻需要在一個 worker 上進行計算，是本地化的。
如上圖所示的例子，如果 worker 1 上有 block A， block B 和 block C，基于 LRU 算出來 block C 是最長時間沒有使用的，就會把 block C淘汰。如果看一下全局的情況，就會發現這樣并不好。因為 block C 在整個叢集中隻有一個副本。把它淘汰之後，如果下面還有人要通路 block C，隻能從遠端拉取資料，就會帶來性能和成本的損失。我們提出做一個全局的淘汰政策。在這種情況下，不應該淘汰 block C，而應該淘汰副本比較多的。在這個例子中，應該淘汰 block A，因為它在其它的節點上仍然有兩個副本，無論是成本還是性能都要更好。
如上圖所示，我們做的工作是在每個 worker 上維護副本資訊。當某一個 worker，比如加了一個副本，或者減了一個副本，首先會向 master 彙報，而 master 會把這個資訊作為心跳傳回值，傳回給其它相關的 worker。其它 worker 就可以知道整個全局副本的實時變化。同時，更新副本資訊。是以當進行 worker 内部的淘汰時，可以知道每一個 worker 在整個全局有多少個副本，就可以設計一些權重。比如仍然使用 LRU，但是會加上副本個數的權重，綜合考量淘汰和替換哪些資料。
經過我們初步的測試，在很多領域，無論是 big data，AI training 中都可以帶來很大的提升。是以不僅僅是優化一台機器上一個 worker 的緩存命中。我們的目标是使得整個叢集的緩存命中率都得到提升。
最後，對全文進行一下總結。首先，在 AI 的訓練場景中，uniform 緩存淘汰算法要比傳統的 LRU、LFU 更好。第二，緩存和遠端的 networking 也是一個需要被配置設定和排程的資源。第三，在進行緩存優化時，不要隻局限在一個作業或者一個 worker 上，應該統攬整個端到端全局的參數，才能使得整個叢集的效率和性能有更好的提升。
以上就是本次分享的内容，謝謝大家。
今日議題推薦
如何基于開源LLM模型快速建構類ChatGPT應用？
⏰ 活動時間：6/27-28 9:00-17:00
☕️ 活動地點：上海·世博中心
點選下面連結報名觀看：
2023年亞馬遜雲科技中國峰會 - 因建構_而可見