面向大模型的存储加速方案设计和实践

作者：Baidu AICLOUD

2023-07-10
北京
本文字数：6607 字
阅读完需：约 22 分钟

这是 AI 大底座系列云智公开课的第三期内容。前两期我的两位同事已经向大家介绍了高性能网络和 GPU 容器虚拟化的相关内容。今天我们把目光聚焦在存储方向，一起来看看面向大模型的存储加速方案的设计和实践。

今天将从以下三个方面来展开这次分享：

介绍大模型全流程对存储带来的全新挑战；
深入大模型全流程各个环节，看一看有哪些具体的存储问题以及对应的解决思路；
分享百度沧海·存储的加速方案及实践经验。

1. 大模型对存储的全新挑战

从过去的经典 AI，到今天人人谈论的大模型，我们看到 AI 模型的参数规模呈现出指数级的爆发增长。一方面，大模型的应用效果开始给大家带来非常大的惊喜，另一方面，也给整个基础设施带来巨大的挑战。

其一，模型规模大，训练时间长。一个 175B 参数的模型，万卡同时训练仍然需要长达 22 天。这就要求基础设施提供超高的性能和超长时间的稳定。

其二，大模型要结合具体应用才能发挥巨大的威力。大家今天谈论大模型，不再只停留在模型本身，更多的关注已经聚焦于结合业务的应用落地。面对互联网级的应用迭代，要求我们具备大规模的敏捷部署能力。

第三，大模型离不开持续更新的海量数据，这就需要与整个数据生态互通，让数据能在各个环节便捷地流动。

在这样的背景下，我们来对大模型全流程做一个拆分，大致可以划分为四个主要的环节。

第一是海量数据的存储和处理，包括采集导入、清洗、转换、标注、共享和长期归档，是后面各环节的基础。这里对存储的要求跟以前的大数据应用具有很大的共性，也带有大模型自身的特点，总结起来主要是生态的互通、高吞吐和大容量。
第二是模型开发，讲究效率为王，包括实验管理、交互式开发和效果评估等。对存储的要求更多集中在 POSIX 兼容性、可靠性和可共享等方面。
第三是模型训练。真正做过大模型训练的朋友一定深有体会，每分每秒都是经费在燃烧。所以时间就是金钱，拒绝等待，拒绝失败。这里的主要场景，一是训练数据的读取，二是为了容错做的 checkpoint 的保存和加载。数据集的部分就是要尽量读得快，减少计算对 I/O 的等待，而 checkpoint 主要要求高吞吐、减少训练中断的时间。
最后是模型推理，需要把训练完的模型快速分发部署到线上，产生业务效果。而这个过程会高频、反复发生，既要求高并发、高吞吐，又要求整个流程尽量简单高效。

从这四个环节的分析，我们总结出大模型对存储的第一个挑战：不同环节对存储提出了复杂多样的需求。

进一步，我们又把大模型的各个环节按照业务流程串联起来。整体上可分为大模型应用、数据系统和 AI 系统三大部分。

其中数据系统的部分在过去的大数据应用中大家都比较熟悉了，最终是让数据与应用产生一个正向反馈的闭环。

而加入 AI 系统以后，面临着更多环节的数据流动。首先经过预处理的数据由数据系统流入 AI 系统被加载训练，训练结果进入模型仓库，再由模型仓库部署到线上提供推理服务，而推理效果的相关数据又会回到数据系统，用于下一步的评估和后续迭代。

因此可以看到，大模型应用、数据系统、AI 系统三大部分实际上构成了一个更大的闭环。其中各环节间的衔接，本质上都是对数据广泛、高效流动的需求。这是大模型对于存储的第二个挑战。

接下来，我们就带着这两个挑战，一起来看大模型各环节具体存储问题的解决思路。

2. 大模型全流程存储问题的解决思路

我们首先来关注第二个挑战，解决海量数据的存储和流动问题，这是解决其它问题的基础。

当我们把数据系统中，数据流入、处理和流出的具体手段稍作展开，就会发现大模型依赖的数据需要与如此广泛的生态频繁交互，基于原来的本地存储或者自建的小规模商用存储已经无法充分利用这些生态的优势。因此数据湖存储已经成为这里事实上的首选方案。

而数据湖存储其实有两个主要的选择。一是以 HDFS 为代表的分布式文件系统，另一个是对象存储。从两类系统的架构来对比，可以观察到两点：

其一是扩展能力。HDFS 采用集中式元数据管理，规模相对受限，而对象存储通常采用分布式平坦元数据，具有更强的水平扩展能力，单集群可支持万亿文件、EB 级规模。在我们对于实践的观察中，曾经有业务基于 HDFS 做大模型训练，受限于扩展能力，不得不将海量小文件进行打包存储，在训练时再下载解压，实际上引入了非常多的流程开销；
其二是成本考量。HDFS 诞生于存算一体的时代，而对象存储天然面向存算分离设计，结合 EC 编码、分级存储等丰富的能力，可以实现大规模数据长期保存的更优成本。

另外，从可用性、吞吐、生态支持等方面来对比，对象存储也具有一定的优势。所以虽然 HDFS 起步较早，但综合来看对象存储才是数据湖存储的更优选择。

尤其是在今天这个云原生的时代，基于对象存储的数据湖更是成为了云上的数据中枢。

带着这个结论回到我们开始的问题，可以看到利用基于对象存储的云原生数据湖就能很好地解决数据系统内各环节的衔接。但是 AI 系统各环节的数据流动和衔接问题，能否也用同一套存储来解决呢？

假如答案是肯定的，那么我们就能得到一个非常统一和简洁的存储架构。最下面是数据湖，上面承接从数据系统到 AI 系统所有环节的需求，数据流动会被大大简化，各环节的衔接运转效率也会得到很大的提升。

而这里的关键就是要接着回答我们最开始提出的第一个挑战，如何用对象存储满足大模型开发、训练、推理等各环节的多样化需求。

首先我们以大家比较关心的模型训练环节为例做一个分析。

对于一个典型的训练来说，可能迭代多轮 epoch。在每个 epoch 内，首先需要对数据集进行随机打散，然后将打散后的数据划分为若干 batch，每读取一个 batch 的数据，进行一次训练迭代。同时会周期性保存 checkpoint 用于故障快速恢复。

如果我们把训练过程进一步展开，可以观察到这样的时序关系。每一轮 epoch 的耗时都是由数据 shuffle、数据读取等待、checkpoint 和真正的训练时间相加得到。因此为了尽量提高训练效率，减少 GPU 的空闲，我们的主要优化就集中在三个思路：

优化 shuffle 过程，尽量将耗时控制在较小比例；
优化读取过程，尽量让每一轮读取数据的耗时小于计算耗时，这样就能让 I/O 时间被计算时间完全隐藏掉；
优化 checkpoint 过程，尽量缩短 checkpoint 耗时，减少训练中断时间。

接下来我们先分析数据的 shuffle 和读取，checkpoint 过程随后再讨论。

如果我们对 shuffle 和读操作进行具体的分析，就会发现，shuffle 是一个纯元数据操作的过程，主要依赖大量文件的 LIST，读取过程则元数据和数据操作都有。

而对于大模型训练用到的 KB 级小文件而言，文件越小，元数据操作的耗时占比也就越高；I/O 越小，延时和 QPS 越超过吞吐成为主要矛盾。因此，对于小文件的 shuffle 和读取，延时和 QPS 是提升性能的关键。

分析完典型的训练过程以后，我们再来看对象存储面对小文件时有哪些影响性能的因素。

其一，元数据方面，对象存储采用分布式 KV 或 NewSQL 维护平坦目录元数据，很好地解决了小文件的规模扩展问题。虽然大部分操作性能很好，但恰恰 shuffle 用到的 LIST 操作性能在部分场景下不够理想，延时偏高；

其二，对象存储协议的整个处理路径较长，对于大文件吞吐为主的场景可以忽略，但对于小文件的频繁读取则会带来不可忽视的延时开销。

不难发现，对象存储的高吞吐能力给小文件读取带来的帮助比较有限，而在延时上也并不具有优势。面对这样的问题，我们有哪些解决思路呢？

我们分为几个方向来总结：

让敌人变弱，也就是通过一些打包格式减少小文件元数据操作的占比。比如 TFRecord、HDF5 等格式已经被大量应用。如果存储能配合这些格式做一些优化，那么效率将得到进一步提升；
让自己变强。硬件上通过燃烧经费购买高性能硬件不需要过多解释。软件上则一般通过架构升级提升系统的扩展能力，缩短 I/O 路径，这里的典型方案就是并行文件系统；
尽可能近身攻击。让存储和计算靠得更近，这里的典型方案就是缓存系统。利用数据集只读的特点，将元数据和数据缓存到计算节点，或者靠近计算节点的地方，也能大幅提高数据的访问效率。

从整体来看，前面我们已经基于对象存储的云原生数据湖解决了海量数据的存储和流动问题。这里我们可以进一步基于并行文件系统或缓存系统的数据存储加速层来弥补对象存储的短板，满足大模型各环节的性能需求。

有了这一套「数据湖 + 加速层」的思路，我们就来详细看看大模型的训练、推理几个具体场景下的问题如何被一一解决。

第一个场景，先来看数据集的读取加速。

假如数据已经进入加速层，那么读性能的问题就能得到很好的解决。剩下需要解决的是数据怎么进入加速层的问题。

在过去我们大量依赖人工和外围工具做不同存储之间的数据流转，带来不少弊端。今天的加速层设计，通常会选择内置自动的数据流转功能。比如在加速层通过 bucket link 建立对象存储某一个 bucket 到加速层的关联，就能自动完成数据在两层之间的流动。这种方式有诸多优点：

同步速度快，各节点可以并发同步数据，大大加快数据流转效率；
同步策略可定制，可以按照场景需求进行增量同步、选择性同步；
能够与调度器集成，实现数据准备与资源调度的高效配合；
避免了人工方式需要处理各类异常，做垃圾回收的问题。

这里就是一个基于自动数据流转实现调度优化，减少训练过程数据读取等待的例子。当两个任务都需要先加载数据然后才能开始训练，通过训练平台的流水线化调度，在一个任务做训练的同时发起下一个任务所需数据的提前加载，就能大大提高计算资源的利用率。

第二个场景，我们接着来看训练中的 checkpoint 部分。

与训练数据以小文件为主不同，大模型单个节点的 checkpoint 通常就能达到几十上百 GB。而多个训练节点同时写，需要恢复时又同时读，对存储提出了很高的吞吐要求。同时一个关键的问题是 checkpoint 期间整个训练是中断的。因此通过提高吞吐，将 checkpoint 耗时控制在尽量小的占比，对于减少训练等待、提高计算资源利用率同样非常重要。

在之前有一种朴素的做法。训练程序先将 checkpoint 写到性能相对容易保证的本地存储，然后再通过外围工具向远端对象存储上传。这里需要考虑一个问题，如果要保证每个 checkpoint 都成功写入，那么总耗时就等于写本地耗时加上传对象存储的耗时，总体等待时间较长。因此实际应用中也有延迟一个 checkpoint 异步上传的方案。但无论如何，都引入了额外的复杂度，需要外围工具来处理各种异常情况。

现在基于数据湖存储加速以后，我们可以有新的做法。训练程序直接将 checkpoint 写入加速层的 Memory 或 NVME SSD，加速层采用流式上传的方式，不必等待 checkpoint 数据全部写完就开始向对象存储上传。此外，加速层也会采用分块上传的办法，充分利用对象存储的后端并发能力。

这个方案看上去比朴素方案有了很大的进步，但是由于需要保证 close 完成时数据已经完整写入对象存储，因此吞吐瓶颈仍然可能出现在上传对象存储的带宽限制上。我们称这种方案为同步方案。

其实对 checkpoint 场景稍加分析就能发现，我们并不一定需要同步方案，因为训练过程的 checkpoint 会周期不断地进行，而发生故障恢复不应该是常态。因此要突破上传对象存储的带宽限制，异步写是一种值得尝试的思路。

这个方案最大的变化，就是对 checkpoint 文件的 close 操作变成了异步，训练程序不用等待数据上传完成，即可恢复训练，剩下的工作全部交给加速层透明完成。此时耗时主要取决于加速层的 Memory 和 NVME SSD 写入能力，可以大大缩短 checkpoint 写入时间。

另外，这里还有一个相关的优化，就是对于最新的一个 checkpoint 采用异步写的同时，让它驻留在加速层的缓存内。当训练需要恢复的时候就能直接从加速层读取，解决了 checkpoint 的快速加载问题。

最后我们来看第三个场景，也就是推理阶段的模型分发部署。大模型的部署通常是百 GB 模型文件的高并发重复读取。由于推理服务规模大，模型迭代更新频繁，我们需要从提高吞吐和缩短流程两个方面考虑如何优化。

首先来看过去常用的基于镜像分发的方案。

训练产出的模型首先需要落到临时存储，完成镜像的制作，包括数据打包、压缩等过程，然后再从临时存储写入持久化的镜像仓库。在推理部署时，再从镜像仓库并发拉取到各推理实例的本地存储，然后进行解压和数据校验。可以看到在这个方案下，吞吐主要取决于镜像仓库底层存储的能力，而流程上在镜像制作和镜像分发两个阶段都需要引入额外的开销。

如果基于数据湖存储加速的方案，整个过程会变得非常简单。

首先从流程上看，训练产出的模型直接写入统一的数据湖存储。这时可以通过对象存储的事件通知机制，让加速层立即感知并提前把模型文件的元数据和数据预加载到靠近推理服务的缓存内。当启动模型部署时，就能直接从缓存读取到数据。

其次从吞吐上看，基于 Memory 或 NVME SSD 提供靠近推理服务的分布式缓存，提供了很强的读吞吐能力，并且多实例重复读取同一份模型数据时不需要消耗大量的对象存储 I/O 和带宽。

另外，加速层可以通过不同的策略满足不同规模的分发加速需求。比如对于几十实例以下的小规模部署，采用单副本即可将所有缓存盘的 I/O 能力均衡地利用起来，满足读性能要求，同时还能大幅节省缓存空间，保证更多模型数据的缓存命中率。而对于数百实例的部署规模，采用多副本即可提高加速层的读吞吐能力。到了数千实例的规模，还能进一步结合客户端的 P2P 加速满足性能需求。

到这里为止，我们分析了模型训练、推理等环节的三个具体场景。可以看到，在引入数据湖存储加速层以后，这些场景的具体问题都一一得到了解决。

因此，「数据湖 + 加速层」能够同时解决最开始提出的两个挑战，为大模型全流程提供了统一的存储解决方案。

同时也能看到，使用「数据湖 + 加速层」这套统一的存储方案，既能获得如同过去使用本地盘的便捷体验，又能充分享受背后数据湖存储带来的扩展性、高性能和繁荣的生态。这正是一直以来大家对于存储的一大梦想。

3. 百度沧海·存储加速方案和实践

接下来我向大家介绍百度沧海·存储加速的解决方案以及具体实践。

这是百度沧海·存储加速方案的产品架构图。底层是我们的对象存储 BOS，提供了大规模、可扩展、低成本的云原生数据湖存储产品，支持丰富的周边生态和便捷的数据流动。中间就是我们的数据湖存储加速层，有两类产品可供选择。一是并行文件系统 PFS，通过独立部署、高性能硬件和网络、全并行软件架构满足极致性能需求。二是数据湖存储加速 RapidFS，通过近计算部署提供更具性价比的分布式缓存加速能力。最上面是我们的 AI 计算，包括异构算力、高速网络、云原生 AI 平台等。