如何借助分布式存储 JuiceFS 加速 AI 模型训练

传统的机器学习模型，数据集比较小，模型的算法也比较简单，使用单机存储，或者本地硬盘就足够了，像 JuiceFS 这样的分布式存储并不是必需品。

随着近几年深度学习的蓬勃发展，越来越多的团队开始遇到了单机存储的瓶颈，分布式存储在 AI 领域的重要性不断凸显。AI 团队通常会面临以下几种问题：

数据集太大

随着数据量和模型规模的增加，单机存储往往无法满足需求。为解决这些问题，就需要使用分布式存储。

历史数据集需要进行全量归档

在某些应用场景每天都会产生大量新的数据集，这些数据集在一段时间后将变为历史数据，需要进行归档。由于这些数据的生成成本较高，因此不能轻易删除，尤其是在自动驾驶领域，如路测车采集的雷达和摄像头数据，这些数据对公司来说是极为宝贵的资产。若要整理这些数据，传统的单机存储显然不足以满足需求，因此需要考虑使用分布式存储。

小文件和非结构化数据太多

针对传统的分布式文件系统，管理大量小文件会带来元数据存储的负担，对于视觉类模型，影响更加明显。解决这一问题的方法是使用对小文件存储友好的分布式存储系统，这样可以保证上层训练任务的高效率，同时也能够方便地管理大量的小文件。

训练框架需要 POSIX 接口

最初算法科学家在做模型调研时，都是基于本地的资源来做研发和数据的访问，但当需要在分布式存储上进行更大规模的训练时，原本的代码一般不会做太多的调整。所以这就要求分布式存储需要支持 POSIX 接口，最大程度上兼容本地开发阶段的代码。

公共数据集需要不同团队共享，也可能需要数据隔离

在某些领域，如计算机视觉，有一些权威的公共数据集，这些数据集需要在公司内部不同的团队间共享。为了方便团队之间的使用，通常会将这些数据整合并打包存储到一个共享存储中，避免不必要的数据复制和冗余。

云上训练的数据 I/O 效率不高

在云上进行模型训练通常是使用对象存储作为底层存储的存算分离架构，由于对象存储的读写性能较差，在训练上会有很大的瓶颈。

本文将会介绍在模型训练中如何使用 JuiceFS，以及优化训练效率的实践。

1.JuiceFS 在模型训练场景中的架构

上图是架构图，分为三个部分：

第一部分：元数据引擎，根据个人选择，可以使用任何数据库，例如 Redis、MySQL 等等，作为元数据引擎。

第二部分：底层数据存储，在云上或私有云中，使用对象存储服务来对接 JuiceFS。

第三部分：JuiceFS 客户端，用户在使用时需要在每个 GPU 和计算节点上挂载 JuiceFS，这样就可以像访问本地硬盘一样访问 JuiceFS 的文件系统。

底层存储依赖于对象存储中的原始数据，同时每个计算节点上还有一些本地缓存，包括元数据和数据缓存。JuiceFS 的设计中，每个计算节点的本地可以有多级缓存。第一级是基于内存的缓存，第二级是基于本地磁盘的缓存，只有在本地缓存没有命中时，才会访问对象存储。

如果进行单机模型训练，在首轮训练时，训练集或数据集通常不会命中缓存。但是从第二轮开始，在缓存资源充足的情况下，几乎不需要访问对象存储，达到加速数据 I/O 的效果。

JuiceFS 读写缓存流程

我们之前进行了一项评测，比较了在访问对象存储时，使用缓存和不使用缓存这两种方式对于训练效率的影响。评测结果表明，这两种方式的性能差别非常大。（点击此处，了解评测结果）

上图展示了 JuiceFS 缓存读写的流程。最上面是应用程序，它相当于是最初发起读请求的起点应用或训练任务。

当应用程序发起读请求后，请求会先进入左侧的内核空间，内核会查看内核页缓存中是否有请求的数据。如果内核页缓存中没有数据，请求会回到用户空间的 JuiceFS 进程。在用户空间，JuiceFS 进程会处理所有的读写请求。

JuiceFS 默认会在内存中维护一个读缓冲区，当请求未能从缓冲区中获取数据时，JuiceFS 会进入块缓存索引，即基于本地磁盘的缓存目录。JuiceFS 默认将文件切割成 4MB 的块并存储，因此缓存的粒度也是 4MB。

举个例子，当访问一个文件的一部分数据时，JuiceFS 只会缓存该部分数据对应的 4MB 块到本地缓存目录中，而不会缓存整个文件。这是 JuiceFS 与其他文件系统或缓存系统的显著差异之一。

Block cache index 用于快速定位本地缓存目录中的文件块。如果找到了，JuiceFS 进程会自动读取本地盘，然后进入内核态，读取完成后返回给 JuiceFS 进程，最后返回给应用程序。

当本地盘数据读取完成后，数据还会进入内核页缓存。这是因为如果没有使用 direct I/O，Linux 系统默认会将数据存储在内核页缓存中。这些内核页缓存都用于加速缓存访问，如果第一个请求直接命中并返回，那么效率是最高的，并且请求不会通过 FUSE 层进入用户态进程。如果没有命中，则会通过 index 查找，如果在节点目录中没有找到 block，则会通过网络请求到达对象存储，然后将数据读回来并原路返回给应用程序。

从对象存储上下载数据时，JuiceFS 会有一个后台异步线程，把读回来的 block 同时写到本地缓存盘里，确保下一次访问同样的 block 时，能直接从本地缓存命中，而不需要再次从对象存储上获取。以上就是 JuiceFS 建立读缓存的流程。

上图有一部分的模块叫 Chunk Cache，chunk 是 JuiceFS 中的一个逻辑概念，每个文件会按照 64MB 大小分为多个 chunk，来提升大文件的读取性能。这部分信息会被缓存到 JuiceFS 进程的内存里，来加速元数据访问的效率。

与数据缓存不同，元数据缓存时间较短，并且为确保强一致性，open 操作默认不缓存。考虑到元数据流量很小，所以对整体的 I/O 性能影响比较小，但是在大量小文件的场景，如果需要频繁访问小文件，元数据的开销也会占到一定的比重。

2.为什么训练太慢以及如何排查？

当使用 JuiceFS 进行训练时，性能是最重要的考虑因素，它直接影响到模型训练的速度。以下是可能影响 JuiceFS 效率的几个方面：

元数据引擎

在处理小文件时，选择不同的元数据引擎（如 Redis、TiKV、MySQL）的性能差别很大。JuiceFS 官网提供了一份对比它们作为元数据引擎的性能文档，平均来说 Redis 会比其他数据库快 3~5 倍。如果发现元数据请求特别慢，建议尝试使用一些性能更好的数据库作为 JuiceFS 的元数据引擎。

对象存储

主要影响数据存储访问的性能和吞吐量。如果在云上使用，通常使用公有云提供的对象存储服务，其性能相对固定。如果使用自建的对象存储，例如使用开源的 Ceph 或 MinIO 组件，可以对组件进行调优以达到更好的性能和吞吐量。

本地磁盘

缓存目录存储的位置对整个读取性能影响很大。在缓存命中率高的情况下，缓存磁盘的 I/O 效率会直接影响整体 I/O 效率。因此需要注意存储类型、存储介质以及磁盘容量等因素，数据集的大小也会对训练效率产生影响。

网络带宽

在第一轮训练完成后，如果数据集不足以在本地完全缓存，网络带宽或网络资源的消耗会影响整体数据访问效率。在云上，不同机型的网卡带宽也有所不同，这也会对数据的访问速度和效率产生影响。

内存

内存的大小会直接影响内核页缓存的大小。当内存足够大时，剩余的空闲内存可以用作 JuiceFS 数据的缓存，进一步加快数据的访问速度。但是，当剩余的空闲内存较少时，数据访问需要通过本地磁盘获取，这会导致访问开销变大。另外，内核态和用户态之间的切换会对性能造成影响，比如系统调用的上下文切换开销等。

如何排查

JuiceFS 提供了许多工具和命令来帮助用户更好地进行性能调优和诊断。在去年的 Office Hours 中，已经对如何在 JuiceFS 中进行性能调优和诊断进行了全面介绍。如果感兴趣，可以在 B 站上观看视频回放。以下是其中几个方法的简要介绍：

工具1 ：juicefs profile 命令

它可以通过分析访问日志来帮助用户更好地优化性能。每个文件系统挂载之后都会有访问日志，但访问日志并不会实时保存，只有在查看访问日志时才会显示出来。相比直接查看原始的访问日志，juicefs profile 命令会进行信息的聚合和类似滑动窗口的数据统计，并按照响应时间从高到低排序，帮助用户优先关注响应时间较慢的请求，进一步分析请求与元数据引擎或对象存储的关系。

工具2：juicefs stats 命令

它从更宏观的角度收集监控数据，并实时展示出来。它可以监控当前挂载点的 CPU 占用、内存占用、内存中的缓冲区占用、FUSE 读写请求、元数据请求以及对象存储的延迟情况等。这些细致的监控指标可以方便用户查看和分析当前模型训练的瓶颈或性能问题出现的可能环节。

JuiceFS 还提供了更底层的信息分析工具，包括 CPU profile 和 heap profile。CPU profile 可以分析 JuiceFS 进程执行速度的瓶颈所在，适用于熟悉源代码的用户。而 heap profile 则主要用于分析内存占用情况，尤其是当 JuiceFS 进程占用大量内存时，需要使用 heap profile 来确定具体哪些函数或数据结构占用了较多内存。

3.一些常见的优化策略

元数据缓存优化

元数据缓存的优化方案主要分为两类：

1）调整内核元数据缓存的超时时间

可以使用 --attr-cache、--entry-cache 和 --dir-entry-cache 参数，这三个参数分别对应不同类型的元数据：attr 表示文件属性（如大小、修改时间、访问时间等），entry 是 Linux 中的概念，表示文件和相关属性，dir-entry 表示目录和其中包含的文件。

这些参数分别控制着元数据缓存的超时时间。为了保证数据访问的一致性，这些参数的默认值只有一秒钟，但在模型训练的场景中，原始数据通常不会被修改，因此可以将这些参数的超时时间设置得更长一些，比如几天到一周等。但需要注意的是，元数据缓存是无法主动失效的，只能等待超时时间到期。

2）优化 JuiceFS 客户端的用户态元数据缓存

默认情况下，在打开文件时会强制请求元数据引擎获取最新的文件属性，以保证强一致性。但由于模型训练的数据通常不会被修改，因此可以打开 --open-cache 参数，并设置一个超时时间，以避免每次打开同一个文件都重复访问元数据引擎。另外可以通过 --open-cache-limit 参数控制缓存的最大文件数，默认值是 10000，即最多缓存最近打开的 10000 个文件的元数据在内存中，可以根据数据集的文件个数进行适当调整。

数据缓存

JuiceFS 的数据缓存分为内核页缓存和本地数据缓存两种，其中内核页缓存无法进行参数调优。因此，在计算节点上应该尽量保留足够的空闲内存，以便 JuiceFS 能够充分利用。如果计算节点上的资源紧张， JuiceFS 就不会将数据缓存到内核中。

而本地数据缓存相对来说用户更加可控，可以根据具体场景调优缓存参数。首先，可以调整缓存的大小（ --cache-size），默认值为 100G，对于大部分场景都足够了。但是对于占用空间特别大的数据集，需要适当调整缓存大小，否则 100G 的缓存空间可能很快被写满，导致 JuiceFS 无法缓存更多数据。配合 --cache-size 参数一起使用的另一个参数是 --free-space-ratio，这个参数用于控制缓存盘的空间空闲比例，默认值是 0.1，即最多使用 90% 的磁盘空间缓存数据。

JuiceFS还支持同时使用多个缓存盘，推荐尽量使用所有可用的盘。数据会通过轮询的方式均匀分布到多个盘中，从而实现负载均衡，同时最大化利用多块盘的存储优势。

缓存预热

为了提高训练效率，可以通过预热缓存来加速训练任务。JuiceFS 支持预热客户端中的元数据缓存和本地数据缓存，通过使用 juicefs warmup 命令可以将缓存提前预热到缓存盘，从而在训练任务开始时直接命中缓存，提高效率。

增大缓冲区大小

缓冲区的大小也会影响读取性能。默认情况下，缓冲区大小为 300MB，但在高吞吐的训练场景下，这可能不够用。可以根据训练节点的内存资源情况来调整缓冲区大小，一般来说，缓冲区越大读取性能越好，但也不要设置过大的值（特别是在限制了最大内存的容器环境中）。需要结合实际负载情况进行调优，找到一个相对合理的缓冲区大小。可以参考前面介绍的 juicefs stats 命令实时观测缓冲区的使用量。

如有帮助的话欢迎关注我们项目 Juicedata/JuiceFS 哟！ (0ᴗ0✿)