如何借助 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]）。

本文的关注点在于最底层的存储层，在保持上层组件