AI 时代我们需要什么样的存储

之前有朋友在询问 LLM 在训练时 Checkpoint 花费的时间太长应该如何解决，正好我在 23 年初和算法团队合作时面临同样的问题。让我们抛开调度问题，先从纯存储角度看一下这个问题的来源。随着 LLM 的爆火，很多公司都躬身入局，加入「百模大战」。不论是训练还是推理，LLM 都对 Infra 提出了新的挑战。传统 DL 包括 NLP 后者图像领域的模型通常规模较小，训练耗时不算长，存储上面临的主要问题都是小文件问题（比如数亿小文件的存储），但是在性能没有太多需求。LLM 动辄 100B 的参数量，使得我们不得不使用更昂贵的硬件比如 A100/H100 来训练模型，这使得我们对于 GPU 的利用率非常敏感。

然而当我们实际开始工作，就会发现 GPU 的利用率总是上不去。有过统计在一个 Epoch（将所有训练样本训练一次的过程）中，高达 70% 的时间都用在数据传送到 GPU 之前。在数据管道（Data Pipeline）的各个阶段，需要花费大量时间在不同系统之间复制数据。实际训练时我们也有过生产的数据，在 20B 的模型规模的训练下，如果我们提前将数据拷贝到本地的 NVMe 磁盘上，一个 epoch 的训练时间为 21s；如果将 dataloader 替换为远端的 s3/hdfs 存储，每个 epoch 的训练时间下降到 40s，GPU 平均利用率不到 50%。可以看到 IO 性能其实对模型训练有很大的影响。

AI 存储的挑战

• 访问接口，也就是应用如何访问分布式文件系统中的数据。分布式文件系统通常都有自己的专有客户端，但是用的很少，目前深度学习领域经常使用的两种协议是基于 HTTP 的 S3 协议以及 POSIX。POSIX 和我们使用本地磁盘没有区别，但是要求分布式文件系统基于 libfuse 实现一套了 VFS 的接口。实际上目前主流的分布式文件系统几乎都提供了 FUSE 接口的支持，但是功能和稳定性一言难尽，实际上能在产生环境稳定运行的并不多。

• 性能，主要是延时和吞吐，大规模训练时还需要考虑 IOPS。在文章开头我们讨论过，延时对于提供 GPU 的利用率至关重要，因此分布式文件系统必须提供低延时的访问，并在整个集群并发读取时保证性能没有太大的下滑。这个问题其实并没有看起来的那么好的解决，比如文件传输链路上的任何中心节点都可能成为整个系统的瓶颈。

• 扩展性，既然分布式文件系统就是为了解决扩展性，为什么还会有扩展性问题呢？这里主要说的还是元数据的扩展性，在 CV 领域数据集通常都是图片和视频，小文件较多，轻松会超过亿级别，而分布式存储集群通常提供了元数据存储规模是严重受到架构限制的。比如 HDFS 采用了中心化的元数据设计且所有元数据保存在内存中，元数据的规模受到单机内存的限制，大概每 GB 能承载在元数据量在 1-2 Million 左右。而对象存储避免了这种扩展问题，将元数据的扩展能力委托给 KV Store，实现了更灵活的扩展。

• 数据管理，这个本身不是一个很难的事情，但是实际操作起来会比较棘手。因为专用的高性能文件系统比较贵，而训练使用的数据集通常比较多，所以会很明显的将存储分为昂贵的高性能训练文件系统和成本低廉的用于持久化的分布式文件系统。我们需要不断的在两个 FS 之间进行数据同步，生命周期管理等操作。在这点上 Alluxio 这类主打 Cache 的系统就比较有优势，云厂商提供的方案里比如百度云的 PFS 也会提供类似的 Bucket Link（通过并行文件系统直接读取对象存储的能力）的功能来解决这个问题。

• 安全，这个是不被讨论但是很重要的服务。假设我们有行业领先的大模型，模型本身就具备相当的商业价值，保护模型本身不被攻击者非法获取也是非常重要。安全这个话题在 AI 时代再次被列中重要的考量当中。

分布式文件系统

分布式文件系统是一种常见的存储方式，目前被大部分公司使用。由于目前网络设备发展迅速，100Gbps ~ 400Gbps 的网卡在训练机器上已经很常见，网络通信的理论速度已经超过了磁盘；此外在容量上，分布式文件系统具备更好的扩展性，不会因为本地磁盘容量的大小限制数据集的大小，灵活性大大增强。

那么在 AI Storage 这个赛道上有哪些解决方案呢？根据笔者目前的观察，大致可以分为两类：一类是传统从 HPC 领域做过来的全闪文件系统以及并行文件系统，由专门的存储提供商或者云厂商提供，这类方案的性能非常好，价格也十分昂贵，每 PB/月的成本大概在千万级别；另外一类方案是传统的分布式文件系统厂商，针对 AI 领域做了优化来提供对应的解决方案，性能比全闪文件系统慢，但是也大致满足需求，每 PB/月的成本大概在百万级别。还有定制硬件来提供加速的公司，比如 weka.io 的估值已经在 7 亿刀左右了。下面大致整理一下对应的公司（没法全部整理，排名不分先后）：

全闪文件系统（All-Flash File System）

全闪文件系统（All-Flash File System）是专门为全闪存（All-Flash）存储阵列设计的文件系统，这种存储阵列完全由固态驱动器（SSD）组成，而不是传统的机械硬盘（HDD）。与传统的文件系统相比，全闪文件系统针对固态驱动器的特性进行了优化，以提高性能和效率，延长固态驱动器的使用寿命。

需要注意的是，全闪方案基本没有开源方案，全部都是商业公司的支持，有些还需要购买专门的硬件，或者需要在公有云依赖特定的硬件支持。目前比较火的是 weka.io 和 VAST，国内的云厂商都有提供对应的解决方案，比如阿里云的 CPFS 以及百度云的 PFS 等。开源的目前发现有 Intel 的 DAOS，但是还没有深入研究。

并行文件系统（Parallel File System）

并行文件系统是一种来自 HPC（高性能计算）领域的存储技术，主要面向科学计算、大数据计算、AI 等领域。并行文件系统面向大规模的存储访问和写入开发，也被顺理成章的引入到 AI 领域。整体来说这个方案没有大问题，但是要注意以下几点：

• 并行文件系统设计的目标是多客户端并行访问大文件，整体设计偏向于顺序读写，需要读写平衡；而目前的 LLM 包括多模态主要的 IO 模式是高并发的随机读，大规模的顺序 checkpoint 写入，和并行文件系统的设计略有差别，达不到最佳性能。

• 并行文件系统的运维管理比较复杂。以 GlusterFS 为例，当节点宕机的时候需要手动的处理数据的再平衡。

• 元数据和容量有限，一般元数据在 1亿，存储容量 1PB，超过这个规模后可能需要非常专业的调优才能继续使用。

传统分布式文件系统（Distribution File System）

这部分其实没什么好讲的，主要特点就是便宜。相比全闪存储和并行文件系统而言，管理运维相对简单一些，但是性能也会差一些。要将这些文件系统用于 AI 训练，也必须进行一些优化才行，主要来说是基于 Cache 的数据本地性优化以及 Offload Kernel 的设计。

我们也可以选择一个分布式文件系统 + Alluxio 的方式来提供 AI 训练的存储方案。和 JuiceFS 类似，Alluxio 提供了分布式的缓存方案以及 FUSE 的 API。在实践中，我们一般选择 CSI 的方式部署这些服务，并尽量将 Cache 节点和训练任务放置在一起。

使用 Alluxio 还有一个特别的好处，我们无需搬运数据，而是直接接入一个我们已有的存储数据集的文件系统即可，大幅降低了数据的管理成本，这是其他方案所不具备的。

CPU Offload

除了文件系统的问题，要做到高性能的数据 IO 和高 GPU 利用率，还有一个重要的问题是优化 CPU 的 workload。以 PyTorch 为例，DataLoader 在读取数据时并不是存粹的 IO ，而是混合了数据预处理，拷贝等操作的。一个数据集被加载到 GPU 之前，既有 IO 处理，也需要 CPU 参与对数据进行一些预处理来满足 GPU 的格式要求。整个 IO 的 Pattern 是并发的随机读，IO 和 CPU 的问题都可能导致数据无法及时送到 GPU 使得 GPU 空闲，利用率低。

GDS（GPU Direct Storage）

我们可以通过将数据预先拷贝到本地的 NVME 磁盘上，在借助 GDS 实现高性能的数据加载。一些分布式文件系统也支持 GDS，比如 BeeGFS 和 CubeFS 具体的效果没有详细进行过测试，有机会了再说。N 厂针对 GDS 还有专门的 SDK: DALI，不知道是不是大力出奇迹的意思。

和并行文件系统一样，RDMA 最早也是应用与 HPC 领域。大数据相关有非常多的基于 RMDA 的论文发表，实际能跑在生产环境的寥寥无几。这是因为 RDMA 需要定制的网卡和交换机的支持，这件事情想想就很贵，而大数据最初诞生就是为了低成本的解决 计算和存储的 scale 的问题。

RDMA 目前的实现主要有三种：IB（InfiniBand）、iWARP 和 RoCE，性能对比如下面的表格。值得一提的是，RoCE 是唯一一个不需要特殊硬件设备的 RDMA 协议，我个人也最看好该技术的发展。越是朴素、廉价的技术越容易得到传播和发展，最终超越昂贵的方案，农村包围城市的故事在技术领域也是不断地在上演。Meta 最新公布的 Llama 3 的训练架构中就有提到，他们同时使用了 RoCE 和 InfiniBand 协议组成的 RDMA 网络，在机器之间提供 400Gpbs 的通信速率。

总结