一、項目背景和緩存策略

首先來分享一下相關背景。

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

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

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

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

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

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

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

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

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

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

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

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

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

不僅是 LRU，像 LFU 等主流的淘汰算法，都存在這樣一個問題。因為整個 AI 訓練對數(shù)據(jù)的訪問是非常均等的。所以，可以采用最簡單的緩存算法，只要緩存一部分數(shù)據(jù)就可以，永遠不用動。在一個作業(yè)來了以后，永遠都只緩存一部分數(shù)據(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)所有后面的訪問，都不會再命中緩存。原因是在進行 LRU 緩存的時候，把它替換出來，但其實在一個 epoch 中已經被訪問一次，這個 epoch 中就永遠不會再被訪問到了。LRU 反倒把它進行緩存了，LRU 不但沒有幫助，反倒是變得更糟糕了。不如使用 uniform，比如下面這種方式。

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

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

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

二、SiloD 框架

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

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

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

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

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

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

對上述介紹進行一下總結：

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

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

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

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

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

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

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

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

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

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