如何借助 JuiceFS 为 AI 模型训练提速 7 倍

原创
09/18 17:58
阅读数 3.5K

背景

海量且优质的数据集是一个好的 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 分布式集群

展开阅读全文
打赏
0
6 收藏
分享
加载中
更多评论
打赏
0 评论
6 收藏
0
分享
返回顶部
顶部