一、項目背景和緩存策略

首先來分享一下相關(guān)背景。

近年來，AI 訓(xùn)練應(yīng)用越來越廣泛。從基礎(chǔ)架構(gòu)角度來看，無論是大數(shù)據(jù)還是 AI 訓(xùn)練集群中，大多使用存儲與計算分離的架構(gòu)。比如很多 GPU 的陣列放到一個很大的計算集群中，另外一個集群是存儲。也可能是使用的一些云存儲，像微軟的 Azure 或者是亞馬遜的 S3 等。

這樣的基礎(chǔ)架構(gòu)的特點(diǎn)是，首先，計算集群中有很多非常昂貴的 GPU，每臺  GPU 往往有一定的本地存儲，比如 SSD 這樣的幾十 TB 的存儲。這樣一個機(jī)器組成的陣列中，往往是用高速網(wǎng)絡(luò)去連接遠(yuǎn)端，比如 Coco、 image net、YouTube 8M 之類的非常大規(guī)模的訓(xùn)練數(shù)據(jù)是以網(wǎng)絡(luò)進(jìn)行連接的。

如上圖所示，數(shù)據(jù)有可能會成為下一個 AI 訓(xùn)練的瓶頸。我們觀察到數(shù)據(jù)集越來越大，隨著 AI 應(yīng)用更加廣泛，也在積累更多的訓(xùn)練數(shù)據(jù)。同時 GPU 賽道是非常卷的。比如 AMD、TPU 等廠商，花費(fèi)了大量精力去優(yōu)化硬件和軟件，使得加速器，類似 GPU、TPU這些硬件越來越快。隨著公司內(nèi)加速器的應(yīng)用非常廣泛之后，集群部署也越來越大。這里的兩個表呈現(xiàn)了關(guān)于數(shù)據(jù)集以及 GPU 速度的一些變化。之前的 K80 到 V100、 P100、 A100，速度是非常迅速的。但是，隨著速度越來越快，GPU 變得越來越昂貴。我們的數(shù)據(jù)，比如 IO 速度能否跟上 GPU 的速度，是一個很大的挑戰(zhàn)。

如上圖所示，在很多大公司的應(yīng)用中，我們觀察到這樣一個現(xiàn)象：在讀取遠(yuǎn)程數(shù)據(jù)的時候，GPU 是空閑的。因為 GPU 是在等待遠(yuǎn)程數(shù)據(jù)讀取，這也就意味著 IO 成為了一個瓶頸，造成了昂貴的 GPU 被浪費(fèi)。有很多工作在進(jìn)行優(yōu)化來緩解這一瓶頸，緩存就是其中很重要的一個優(yōu)化方向。這里介紹兩種方式。

第一種，在很多應(yīng)用場景中，尤其是以 K8s 加 Docker 這樣的基礎(chǔ) AI 訓(xùn)練架構(gòu)中，用了很多本地磁盤。前文中提到 GPU 機(jī)器是有一定的本地存儲的，可以用本地磁盤去做一些緩存，把數(shù)據(jù)先緩存起來。

啟動了一個 GPU 的 Docker 之后，不是馬上啟動 GPU 的 AI 訓(xùn)練，而是先去下載數(shù)據(jù)，把數(shù)據(jù)從遠(yuǎn)端下載到 Docker 內(nèi)部，也可以是掛載等方式。下載到 Docker 內(nèi)部之后再開始訓(xùn)練。這樣盡可能的把后邊的訓(xùn)練的數(shù)據(jù)讀取都變成本地的數(shù)據(jù)讀取。本地 IO 的性能目前來看是足夠支撐 GPU 的訓(xùn)練的。VLDB 2020 上面，有一篇 paper，CoorDL，是基于 DALI 進(jìn)行數(shù)據(jù)緩存。

這一方式也帶來了很多問題。首先，本地的空間是有限的，意味著緩存的數(shù)據(jù)也是有限的，當(dāng)數(shù)據(jù)集越來越大的時候，很難緩存到所有數(shù)據(jù)。另外，AI 場景與大數(shù)據(jù)場景有一個很大的區(qū)別是，AI 場景中的數(shù)據(jù)集是比較有限的。不像大數(shù)據(jù)場景中有很多的表，有各種各樣的業(yè)務(wù)，每個業(yè)務(wù)的數(shù)據(jù)表的內(nèi)容差距是非常大的。在 AI 場景中，數(shù)據(jù)集的規(guī)模、數(shù)據(jù)集的數(shù)量遠(yuǎn)遠(yuǎn)小于大數(shù)據(jù)場景。所以常常會發(fā)現(xiàn)，公司中提交的任務(wù)很多都是讀取同一個數(shù)據(jù)。如果每個人下載數(shù)據(jù)到自己本地，其實(shí)是不能共享的，會有非常多份數(shù)據(jù)被重復(fù)存儲到本地機(jī)器上。這種方式顯然存在很多問題，也不夠高效。

接下來介紹第二種方式。既然本地的存儲不太好，那么，是否可以使用像 Alluxio 這樣一個分布式緩存來緩解剛才的問題，分布式緩存有非常大的容量來裝載數(shù)據(jù)。另外，Alluxio 作為一個分布式緩存，很容易進(jìn)行共享。數(shù)據(jù)下載到 Alluxio 中，其他的客戶端，也可以從緩存中讀取這份數(shù)據(jù)。這樣看來，使用 Alluxio 可以很容易地解決上面提到的問題，為 AI 訓(xùn)練性能帶來很大的提升。微軟印度研究院在 FAST2020 發(fā)表的名為 Quiver 的一篇論文，就提到了這樣的解決思路。但是我們分析發(fā)現(xiàn)，這樣一個看似完美的分配方案，還是比較靜態(tài)的，并不高效。同時，采用什么樣的 cache 淘汰算法，也是一個很值得討論的問題。

如上圖所示，是使用 Alluxio 作為 AI 訓(xùn)練的緩存的一個應(yīng)用。使用 K8s 做整個集群任務(wù)的調(diào)度和對 GPU、CPU、內(nèi)存等資源的管理。當(dāng)有用戶提交一個任務(wù)到 K8s 時，K8s 首先會做一個插件，通知 Alluxio 的 master，讓它去下載這部分?jǐn)?shù)據(jù)。也就是先進(jìn)行一些熱身，把作業(yè)可能需要的任務(wù)，盡量先緩存一些。當(dāng)然不一定非得緩存完，因為Alluxio 是有多少數(shù)據(jù)，就使用多少數(shù)據(jù)。剩下的，如果還沒有來得及緩存，就從遠(yuǎn)端讀取。另外，Alluxio master 得到這樣的命令之后，就可以讓調(diào)度它的 worker 去遠(yuǎn)端?？赡苁窃拼鎯?，也可能是 Hadoop 集群把數(shù)據(jù)下載下來。這個時候，K8s 也會把作業(yè)調(diào)度到 GPU 集群中。比如上圖中，在這樣一個集群中，它選擇第一個節(jié)點(diǎn)和第三個節(jié)點(diǎn)啟動訓(xùn)練任務(wù)。啟動訓(xùn)練任務(wù)之后，需要進(jìn)行數(shù)據(jù)的讀取。在現(xiàn)在主流的像 PyTorch、Tensorflow  等框架中，也內(nèi)置了 Prefetch，也就是會進(jìn)行數(shù)據(jù)預(yù)讀取。它會讀取已經(jīng)提前緩存的 Alluxio 中的緩存數(shù)據(jù)，為訓(xùn)練數(shù)據(jù) IO 提供支持。當(dāng)然，如果發(fā)現(xiàn)有一些數(shù)據(jù)是沒有讀到的，Alluxio 也可以通過遠(yuǎn)端進(jìn)行讀取。Alluxio 作為一個統(tǒng)一的接口是非常好的。同時它也可以進(jìn)行數(shù)據(jù)的跨作業(yè)間的共享。

如上圖所示，比如又有一個人提交了同樣數(shù)據(jù)的另一個作業(yè)，消耗的是同一個數(shù)據(jù)集，這個時候，當(dāng)提交作業(yè)到 K8s 的時候，Alluxio 就知道已經(jīng)有這部分?jǐn)?shù)據(jù)了。如果 Alluxio 想做的更好，甚至是可以知道，數(shù)據(jù)即將會被調(diào)度到哪臺機(jī)器上。比如這個時候調(diào)度到 node 1、node 3 和 node 4 上。node 4 的數(shù)據(jù)，甚至可以做一些副本進(jìn)行拷貝。這樣所有的數(shù)據(jù)，即使是 Alluxio 內(nèi)部，都不用跨機(jī)器讀，都是本地的讀取。所以看起來 Alluxio 對 AI 訓(xùn)練中的 IO 問題有了很大的緩解和優(yōu)化。但是如果仔細(xì)觀察，就會發(fā)現(xiàn)兩個問題。

第一個問題就是緩存的淘汰算法非常低效，因為在 AI 場景中，訪問數(shù)據(jù)的模式跟以往有很大區(qū)別。第二個問題是，緩存作為一種資源，與帶寬（即遠(yuǎn)程存儲的讀取速度）是一個對立的關(guān)系。如果緩存大，那么從遠(yuǎn)端讀取數(shù)據(jù)的機(jī)會就小。如果緩存很小，則很多數(shù)據(jù)都得從遠(yuǎn)端讀取。如何很好地調(diào)度分配這些資源也是一個需要考慮的問題。

在討論緩存的淘汰算法之前，先來看一下 AI 訓(xùn)練中數(shù)據(jù)訪問的過程。在 AI 訓(xùn)練中，會分為很多個 epoch，不斷迭代地去訓(xùn)練。每一個訓(xùn)練 epoch，都會讀取每一條數(shù)據(jù)，并且僅讀一次。為了防止訓(xùn)練的過擬合，在每一次 epoch 結(jié)束之后，下一個 epoch 的時候，讀取順序會變化，會進(jìn)行一個 shuffle。也就是每次每個 epoch 都會把所有數(shù)據(jù)都讀取一次，但是順序卻不一樣。

Alluxio 中默認(rèn)的 LRU 淘汰算法，顯然不能很好地應(yīng)用到AI訓(xùn)練場景中。因為 LRU 是利用緩存的本地性。本地性分為兩方面，首先是時間本地性，也就是現(xiàn)在訪問的數(shù)據(jù)，馬上可能還會即將訪問。這一點(diǎn)，在 AI 訓(xùn)練中并不存在。因為現(xiàn)在訪問的數(shù)據(jù)，在下一輪的時候才會訪問，而且下一輪的時候都會訪問。沒有一個特殊的概率，一定是比其他數(shù)據(jù)更容易被訪問。另一方面是數(shù)據(jù)本地性，還有空間本地性。也就是，為什么 Alluxio 用比較大的 block 緩存數(shù)據(jù)，是因為某條數(shù)據(jù)讀取了，可能周圍的數(shù)據(jù)也會被讀取。比如大數(shù)據(jù)場景中，OLAP 的應(yīng)用，經(jīng)常會進(jìn)行表的掃描，意味著周圍的數(shù)據(jù)馬上也會被訪問。但是在 AI 訓(xùn)練場景中是不能應(yīng)用的。因為每次都會 shuffle，每次讀取的順序都是不一樣的。因此 LRU 這種淘汰算法并不適用于 AI 訓(xùn)練場景。

不僅是 LRU，像 LFU 等主流的淘汰算法，都存在這樣一個問題。因為整個 AI 訓(xùn)練對數(shù)據(jù)的訪問是非常均等的。所以，可以采用最簡單的緩存算法，只要緩存一部分?jǐn)?shù)據(jù)就可以，永遠(yuǎn)不用動。在一個作業(yè)來了以后，永遠(yuǎn)都只緩存一部分?jǐn)?shù)據(jù)。永遠(yuǎn)都不要淘汰它。不需要任何的淘汰算法。這可能是目前最好的淘汰機(jī)制。

如上圖中的例子。上面是 LRU 算法，下面是均等方法。在開始只能緩存兩條數(shù)據(jù)。我們把問題簡單一些，它的容量只有兩條，緩存 D 和 B 這兩條數(shù)據(jù)，中間就是訪問的序列。比如命中第一個訪問的是 B，如果是 LRU，B 存在的緩存中命中了。下一條訪問的是 C，C 并不在 D 和 B，LRU 的緩存中，所以基于 LRU 策略，會把 D 替換掉，C 保留下來。也就是這個時候緩存是 C 和 B。下一個訪問的是 A，A 也不在 C 和 B 中。所以會把B 淘汰掉，換成 C 和 A。下一個就是 D，D 也不在緩存中，所以換成 D 和 A。以此類推，會發(fā)現(xiàn)所有后面的訪問，都不會再命中緩存。原因是在進(jìn)行 LRU 緩存的時候，把它替換出來，但其實(shí)在一個 epoch 中已經(jīng)被訪問一次，這個 epoch 中就永遠(yuǎn)不會再被訪問到了。LRU 反倒把它進(jìn)行緩存了，LRU 不但沒有幫助，反倒是變得更糟糕了。不如使用 uniform，比如下面這種方式。

下面這種 uniform 的方式，永遠(yuǎn)在緩存中緩存 D 和 B，永遠(yuǎn)不做任何的替換。在這樣情況下，你會發(fā)現(xiàn)至少有 50% 的命中率。所以可以看到，緩存的算法不用搞得很復(fù)雜，只要使用 uniform 就可以了，不要使用 LRU、LFU 這類算法。

對于第二個問題，也就是關(guān)于緩存和遠(yuǎn)程帶寬之間關(guān)系的問題。現(xiàn)在所有主流的 AI 框架中都內(nèi)置了數(shù)據(jù)預(yù)讀，防止 GPU 等待數(shù)據(jù)。所以當(dāng) GPU 做訓(xùn)練的時候，其實(shí)是觸發(fā)了 CPU 預(yù)取下一輪可能用到的數(shù)據(jù)。這樣可以充分利用 GPU 的算力。但當(dāng)遠(yuǎn)程存儲的 IO 成為瓶頸的時候，就意味著 GPU 要等待 CPU 了。所以 GPU 會有很多的空閑時間，造成了資源的浪費(fèi)。希望可以有一個比較好的調(diào)度管理方式，緩解 IO 的問題。

緩存和遠(yuǎn)程 IO 對整個作業(yè)的吞吐是有很大影響的。所以除了 GPU、CPU 和內(nèi)存，緩存和網(wǎng)絡(luò)也是需要調(diào)度的。在以往大數(shù)據(jù)的發(fā)展過程中，像 Hadoop、yarn、my source、K8s 等，主要都是調(diào)度 CPU、內(nèi)存、GPU。對于網(wǎng)絡(luò)，尤其對于緩存的控制都不是很好。所以，我們認(rèn)為，在 AI 場景中，需要很好的調(diào)度和分配它們，來達(dá)到整個集群的最優(yōu)。

二、SiloD 框架

在 EuroSys 2023 發(fā)表了這樣一篇文章，它是一個統(tǒng)一的框架，來調(diào)度計算資源和存儲資源。

整體架構(gòu)如上圖所示。左下角是集群中的 CPU 和 GPU 硬件計算資源，以及存儲資源，如 NFS、云存儲 HDFS 等。在上層有一些 AI 的訓(xùn)練框架 TensorFlow、PyTorch 等。我們認(rèn)為需要加入一個統(tǒng)一管理和分配計算和存儲資源的插件，也就是我們提出的 SiloD。

如上圖所示，一個作業(yè)可以達(dá)到什么樣的吞吐和性能，是由 GPU 和 IO 的最小值決定的。使用多少個遠(yuǎn)程 IO，就會使用多少遠(yuǎn)端的 networking?？梢酝ㄟ^這樣一個公式算出訪問速度。作業(yè)速度乘以緩存未命中率，也就是（1-c/d）。其中 c 就是緩存的大小，d 就是數(shù)據(jù)集。這也就意味著數(shù)據(jù)只考慮 IO 可能成為瓶頸的時候，大概的吞吐量是等于（b/(1-c/d)），b 就是遠(yuǎn)端的帶寬。結(jié)合以上三個公式，可以推出右邊的公式，也就是一個作業(yè)最終想達(dá)到什么樣的性能，可以這樣通過公式去計算沒有 IO 瓶頸時的性能，和有 IO 瓶頸時的性能，取二者中的最小值。

得到上面的公式之后，把它微分一下，就可以得到緩存的有效性，或者叫做緩存效率。即雖然作業(yè)很多，但在分配緩存的時候不能一視同仁。每一個作業(yè)，基于數(shù)據(jù)集的不同，速度的不同，緩存分配多少是很有講究的。這里舉一個例子，就以這個公式為例，如果發(fā)現(xiàn)一個作業(yè)，速度非常快，訓(xùn)練起來非常快，同時數(shù)據(jù)集很小，這時候就意味著分配更大的緩存，收益會更大。

基于以上觀察，可以使用 SiloD，進(jìn)行緩存和網(wǎng)絡(luò)的分配。而且緩存的大小，是針對每個作業(yè)的速度，以及數(shù)據(jù)集整個的大小來進(jìn)行分配的。網(wǎng)絡(luò)也是如此。所以整個架構(gòu)是這樣的：除了主流的像 K8s 等作業(yè)調(diào)度之外，還有數(shù)據(jù)管理。在圖左邊，比如緩存的管理，要統(tǒng)計或者監(jiān)控分配整個集群中緩存的大小，每個作業(yè)緩存的大小，以及每個作業(yè)使用到的遠(yuǎn)程 IO 的大小。底下的作業(yè)，和 Alluxio 方式很像，都可以都使用 API 進(jìn)行數(shù)據(jù)的訓(xùn)練。每個 worker 上使用緩存對于本地的 job 進(jìn)行緩存支持。當(dāng)然它也可以在一個集群中跨節(jié)點(diǎn)，也可以進(jìn)行共享。

經(jīng)過初步測試和實(shí)驗，發(fā)現(xiàn)這樣一個分配方式可以使整個集群的使用率和吞吐量都得到非常明顯的提升，最高可以達(dá)到 8 倍的性能上的提升。可以很明顯的緩解作業(yè)等待、GPU 空閑的狀態(tài)。

對上述介紹進(jìn)行一下總結(jié)：

第一，在 AI 或者深度學(xué)習(xí)訓(xùn)練場景中，傳統(tǒng)的 LRU、LFU 等緩存策略并不適合，不如直接使用 uniform。

第二，緩存和遠(yuǎn)程帶寬，是一對伙伴，對整體性能起到了非常大的作用。

第三，像 K8s、yarn 等主流調(diào)度框架，可以很容易繼承到 SiloD。

最后，我們在 paper 中做了一些實(shí)驗，不同的調(diào)度策略，都可以帶來很明顯的吞吐量的提升。

三、分布式緩存策略以及副本管理

我們還做了一些開源的工作。分布式緩存策略以及副本管理這項工作，已經(jīng)提交給社區(qū)，現(xiàn)在處于 PR 階段。Alluxio master 主要做 Meta 的管理和整個 worker 集群的管理。真正緩存數(shù)據(jù)的是 worker。上面有很多以 block 為單位的塊兒去緩存數(shù)據(jù)。存在的一個問題是，現(xiàn)階段的緩存策略都是單個 worker 的，worker 內(nèi)部的每個數(shù)據(jù)在進(jìn)行是否淘汰的計算時，只需要在一個 worker 上進(jìn)行計算，是本地化的。

如上圖所示的例子，如果 worker 1 上有 block A， block B 和 block C，基于 LRU 算出來 block C 是最長時間沒有使用的，就會把 block C淘汰。如果看一下全局的情況，就會發(fā)現(xiàn)這樣并不好。因為 block C 在整個集群中只有一個副本。把它淘汰之后，如果下面還有人要訪問 block C，只能從遠(yuǎn)端拉取數(shù)據(jù)，就會帶來性能和成本的損失。我們提出做一個全局的淘汰策略。在這種情況下，不應(yīng)該淘汰 block C，而應(yīng)該淘汰副本比較多的。在這個例子中，應(yīng)該淘汰 block A，因為它在其它的節(jié)點(diǎn)上仍然有兩個副本，無論是成本還是性能都要更好。

如上圖所示，我們做的工作是在每個 worker 上維護(hù)副本信息。當(dāng)某一個 worker，比如加了一個副本，或者減了一個副本，首先會向 master 匯報，而 master 會把這個信息作為心跳返回值，返回給其它相關(guān)的 worker。其它 worker 就可以知道整個全局副本的實(shí)時變化。同時，更新副本信息。所以當(dāng)進(jìn)行 worker 內(nèi)部的淘汰時，可以知道每一個 worker 在整個全局有多少個副本，就可以設(shè)計一些權(quán)重。比如仍然使用 LRU，但是會加上副本個數(shù)的權(quán)重，綜合考量淘汰和替換哪些數(shù)據(jù)。

經(jīng)過我們初步的測試，在很多領(lǐng)域，無論是 big data，AI training 中都可以帶來很大的提升。所以不僅僅是優(yōu)化一臺機(jī)器上一個 worker 的緩存命中。我們的目標(biāo)是使得整個集群的緩存命中率都得到提升。 

最后，對全文進(jìn)行一下總結(jié)。首先，在 AI 的訓(xùn)練場景中，uniform 緩存淘汰算法要比傳統(tǒng)的 LRU、LFU 更好。第二，緩存和遠(yuǎn)端的 networking 也是一個需要被分配和調(diào)度的資源。第三，在進(jìn)行緩存優(yōu)化時，不要只局限在一個作業(yè)或者一個 worker 上，應(yīng)該統(tǒng)攬整個端到端全局的參數(shù)，才能使得整個集群的效率和性能有更好的提升。