关于后端:如何借助-JuiceFS-为-AI-模型训练提速-7-倍

背景

海量且优质的数据集是一个好的 AI 模型的基石之一，如何存储、治理这些数据集，以及在模型训练时晋升 I/O 效率始终都是 AI 平台工程师和算法科学家特地关注的事件。不论是单机训练还是分布式训练，I/O 的性能都会显著影响整体 pipeline 的效率，甚至是最终的模型品质。

咱们也逐步看到容器化成为 AI 训练的趋势，利用容器能够疾速弹性伸缩的特点，联合私有云的资源池，可能最大化资源利用率，为企业大大节约老本。因而也就诞生了相似 Kubeflow 和 Volcano 这样的开源组件，帮忙用户在 Kubernetes 上治理 AI 工作。Kubernetes 自 1.15 开始新增了 Scheduling Framework，社区也基于这个新的调度框架优化了很多针对 AI 训练场景的问题。后面提到的训练数据管理问题在 Kubernetes 上仍然存在，甚至放大了这个需要，因为计算不再是在固定的几台机器上进行，数据须要智能地追随计算「流动」（或者反过来）。

最初，不论是算法科学家日常试验，还是正式训练模型，POSIX 接口仍然是一个很强烈的需要，尽管支流的框架或者算法库根本都反对对象存储接口但 POSIX 依然是「第一公民」。一些操作系统的高级个性（如 page cache）也是只有 POSIX 接口才具备的。

AI 平台整体架构

下面是一个常见的 AI 平台架构图。存储系统目前应用比拟多的就是对象存储和 HDFS，这里之所以还会用到 HDFS 有多种起因，比方平台部署在机房没有对象存储，训练数据集预处理是在大数据平台等。计算资源混合了 CPU 实例和 GPU 实例，和大数据平台不一样的中央在于，AI 平台的资源天生就是异构的，因而怎么正当高效利用这些异构资源始终是个业界难题。调度器后面曾经介绍到，Kubernetes 是目前支流的组件，联合各种 Job Operator、Volcano、调度插件能够最大水平上施展 Kubernetes 的能力。Pipeline 是很重要的一个局部，AI 工作并不只是由模型训练这一个步骤组成，还包含数据预处理、特色工程、模型验证、模型评估、模型上线等多个环节，因而 Pipeline 治理也是十分重要的。最初就是算法科学家接触最多的深度学习框架，这些框架目前都有本人的应用群体，很多模型优化会基于某种框架进行（比方 TensorFlow 的 XLA），但也有和框架无关的（比方 TVM[4]）。

本文的关注点在于最底层的存储层，在放弃下层组件不变的状况下，如何优化存储层的 I/O 效率。这部分包含但不限于数据缓存、预读、并发读、调度优化等策略，JuiceFS 便是这样一个存储层的加强组件，可能大幅晋升 I/O 效率，上面会具体介绍。

JuiceFS 简介

JuiceFS 是一个面向云原生环境设计的高性能开源分布式文件系统，齐全兼容 POSIX、HDFS、S3 接口，实用于大数据、AI 模型训练、Kubernetes 共享存储、海量数据归档治理等场景。

当通过 JuiceFS 客户端读取数据时，这些数据将会智能地缓存到利用配置的本地缓存门路（可能是内存，也可能是磁盘），同时元数据也会缓存到客户端节点本地内存中。对于 AI 模型训练场景来说，第一个 epoch 实现之后后续的计算都能够间接从缓存中获取训练数据，极大地晋升了训练效率。

JuiceFS 也具备预读、并发读取数据的能力，保障每个 mini-batch 的生成效率，提前准备好数据。

此外 JuiceFS 还提供规范的 Kubernetes CSI Driver，利用能够把 JuiceFS 文件系统作为一个共享的 Persistent Volume（PV）同时挂载到多个容器中。

得益于以上个性的反对，算法科学家能够很轻松地治理训练数据，就像拜访本地存储一样，无需批改框架进行专门的适配，训练效率也能失去肯定的保障。

测试计划

为了验证应用 JuiceFS 当前模型训练的成果，咱们选取常见的 ResNet50 模型以及 ImageNet 数据集，训练任务应用了 DLPerf[5] 我的项目提供的脚本，对应的深度学习框架是 PyTorch。训练节点配置了 8 块 NVIDIA A100 显卡。

作为比照，咱们将私有云上的对象存储作为基准线（通过类 S3FS 的形式进行拜访），同时和开源我的项目 Alluxio 进行比拟，别离测试了 1 机 1 卡、1 机 4 卡、1 机 8 卡不同配置下的训练效率（即每秒解决的样本数）。

不论是 JuiceFS 还是 Alluxio，训练数据集都提前预热到了内存中，数据粗放占用 160G 空间。JuiceFS 提供了 warmup 子命令 [6] 能够很不便地进行数据集的缓存预热，只需指定须要预热的目录或者文件列表即可。

测试方法是每种配置都跑多轮训练，每轮只跑 1 个 epoch，将每轮的统计后果汇总，在排除一些可能的异样数据当前，计算得出整体的训练效率。

JuiceFS 配置选项阐明

AI 模型训练场景的 I/O 模式是典型的只读模式，即只会对数据集产生读申请，不会批改数据。因而为了最大化 I/O 效率，能够适当调整一些配置选项（如缓存相干配置），上面具体介绍几个重要的 JuiceFS 配置选项。

元数据缓存

在内核中能够缓存三种元数据：属性（attribute)、文件项（entry）和目录项（direntry），它们能够通过如下三个选项管制缓存工夫：

--attr-cache value       属性缓存过期工夫；单位为秒 (默认: 1)
--entry-cache value      文件项缓存过期工夫；单位为秒 (默认: 1)
--dir-entry-cache value  目录项缓存过期工夫；单位为秒 (默认: 1)

默认元数据在内核中只缓存 1 秒钟，能够依据训练时长适当增大缓存工夫，如 2 小时（7200 秒）。

当关上一个文件时（即 open() 申请），为了保障一致性[7]，JuiceFS 默认都会申请元数据引擎以获取最新的元信息。因为数据集都是只读的，因而能够适当调整解决策略，设置查看文件是否更新的间隔时间，如果工夫没有达到设定的值，则不须要拜访元数据引擎，能够大幅晋升关上文件的性能。相干配置选项是：

--open-cache value        关上的文件的缓存过期工夫（0 代表敞开这个个性）；单位为秒 (默认: 0)

数据缓存

对于曾经读过的文件，内核会把它的内容主动缓存下来，下次再关上的时候，如果文件没有被更新（即 mtime 没有更新），就能够间接从内核中的缓存（page cache）读取得最好的性能。因而当第一个 epoch 运行结束，如果计算节点的内存足够，那大部分数据集可能都曾经缓存到 page cache 中，这样之后的 epoch 将能够不须要通过 JuiceFS 读取数据，性能也能大幅晋升。这个个性曾经默认在 0.15.2 及以上版本的 JuiceFS 中开启，不须要做任何配置。

除了内核中的数据缓存，JuiceFS 还反对将数据缓存到本地文件系统中，能够是基于硬盘、SSD 或者内存的任意本地文件系统。本地缓存能够通过以下选项来调整：

--cache-dir value         本地缓存目录门路；应用冒号隔离多个门路 (默认: "$HOME/.juicefs/cache" 或 "/var/jfsCache")
--cache-size value        缓存对象的总大小；单位为 MiB (默认: 1024)
--free-space-ratio value  最小残余空间比例 (默认: 0.1)
--cache-partial-only      仅缓存随机小块读 (默认: false)

例如要将数据缓存到内存中有两种形式，一种是将 --cache-dir 设置为 memory，另一种是将其设置为 /dev/shm。这两种形式的区别是前者在从新挂载 JuiceFS 文件系统之后缓存数据就清空了，而后者还会保留，性能上两者没有太大差异。上面是将数据缓存到 /dev/shm/jfscache 并且限定最多应用 300GiB 内存的示例：

--cache-dir /dev/shm/jfscache --cache-size 307200

JuiceFS 也反对将数据缓存到多个门路，默认会采纳轮询的形式写入缓存数据，多个门路通过冒号分隔，例如：

--cache-dir /data1:/data2:/data3

Alluxio 配置选项阐明

Alluxio 的所有组件（如 master、worker、FUSE）都是部署在同一个节点，应用的版本是 2.5.0-2。具体配置如下：

| 配置项 | 设定值 |
| —— | —— |
| alluxio.master.journal.type | UFS |
| alluxio.user.block.size.bytes.default | 32MB |
| alluxio.user.local.reader.chunk.size.bytes| 32MB |
| alluxio.user.metadata.cache.enabled| true |
| alluxio.user.metadata.cache.expiration.time | 2day |
| alluxio.user.streaming.reader.chunk.size.bytes | 32MB |
| alluxio.worker.network.reader.buffer.size | 128MB |

此外 Alluxio FUSE 启动时指定的挂载选项是：kernel_cache,ro,max_read=131072,attr_timeout=7200,entry_timeout=7200,nonempty。

测试后果

测试后果蕴含两个场景，一种应用了内核的 page cache，另一种没有应用。后面提到测试方法是每种配置跑多轮训练，当跑完第一轮当前，后续的测试都有可能间接从 page cache 中读取数据。因而咱们设计了第二种场景，来测试没有 page cache 时的训练效率（比方模型训练的第一个 epoch），这种场景能更实在反映底层存储系统的理论性能。

对于第一种场景，JuiceFS 不须要额定配置即可无效利用内核的 page cache，然而对象存储和 Alluxio 的默认配置都不反对这个个性，须要独自进行设置。

须要特地留神的是，咱们在测试对象存储的过程中已经尝试过开启 S3FS 的本地缓存个性，心愿达到相似 JuiceFS 和 Alluxio 的缓存成果。然而理论测试时发现即便曾经全量预热缓存，以及无论用多少块显卡，1 个 epoch 都无奈在 1 天内跑完，甚至比没有缓存时更慢。因而以下测试后果中的「对象存储」未蕴含开启本地缓存当前的数据。

下图是两个场景的测试后果（「w/o PC」示意没有 page cache）：

得益于元数据缓存和数据缓存，能够看到不论是在哪种场景下，JuiceFS 相比对象存储均匀都能达到 4 倍以上的性能晋升，最多能有靠近 7 倍的性能差距。同时因为对象存储的拜访形式没有无效利用内核的 page cache，因而它在这两种场景的性能差距不大。另外在残缺的端到端模型训练测试中，因为对象存储的训练效率太低，跑到指定模型精度所需工夫过长，在生产环境中根本属于不可用状态。

比照 Alluxio，在有 page cache 的第一种场景中与 JuiceFS 差异不大。在没有 page cache 只有内存缓存的第二种场景中 JuiceFS 均匀晋升 20% 左右的性能，特地是在 1 机 8 卡的配置下，差距进一步加大，达到了 43% 左右的性能差别。Alluxio 在 1 机 8 卡配置下的性能相比 1 机 4 卡没有晋升，没法充分利用多卡的计算能力。

GPU 资源是一种比拟低廉的资源，因而 I/O 效率的差别也能间接体现到计算资源的老本上，越是能高效利用计算资源才越能整体升高 TCO。

总结及瞻望

本文介绍了在 AI 模型训练中如何充分利用 JuiceFS 的个性来为训练提速，相比间接从对象存储读取数据集，通过 JuiceFS 能够带来最多 7 倍的性能晋升。在多卡训练的场景上也能放弃肯定的线性减速比，为分布式训练奠定了根底。

将来 JuiceFS 还会在 AI 场景摸索更多方向，例如进一步晋升 I/O 效率、海量小文件存储、数据与计算的亲和性、与 Job Operator 的联合、与 Kubernetes 调度框架或社区调度器的联合等等。欢送大家积极参与到 JuiceFS 开源社区中来，独特建设云原生 AI 场景的存储基石。

举荐浏览
如何利用 JuiceFS 的性能工具做剖析和调优
基于 JuiceFS 搭建 Milvus 分布式集群