一、項目背景和緩存策略

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

近年來，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 場景中，需要很好的調度和分配它們，來達到整個集群的最優。

二、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 中做了一些實驗，不同的調度策略，都可以帶來很明顯的吞吐量的提升。

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

我們還做了一些開源的工作。分布式緩存策略以及副本管理這項工作，已經提交給社區，現在處于 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 上，應該統攬整個端到端全局的參數，才能使得整個集群的效率和性能有更好的提升。