面向大规模深度学习训练的缓存优化实践

一、项目背景和缓存策略

首先来分享一下相关背景。

近年来，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 上，应该统揽整个端到端全局的参数，才能使得整个集群的效率和性能有更好的提升。

以上是面向大规模深度学习训练的缓存优化实践的详细内容。更多信息请关注PHP中文网其他相关文章！