大模型存储选型 & JuiceFS 在关键环节性能详解

从去年开始,LLM大语言模型领域发展迅速、如 LLaMA、ChatGLM、Baichuan、Qwen 和 yi-model 等基础模型(Foundation Models)的数量显著增加。众多企业也开始基于这些基础模型做 post-training 的相关工作,以开发特定垂直领域的模型实现应用落地。

AI 模型的参数规模呈指数级增长,出现了越来越多的千亿甚至万亿参数的通用大模型。例如,最新的 llama3 模型就提供了 405B、70B 和 8B 三个版本,模型参数量不断刷新,给企业存储的成本与规划带来很大挑战。除了模型规模的增长,大模型开发的复杂流程和高效的数据流转也对存储系统提出了更高的要求,这直接影响到整个流程的效率。

另外,大规模训练和推理对于GPU算力的需求,导致算力的供给越来越多的转变为多云、多region 的模式,在这种情况下,又给用户带来了数据集、模型在多区域间进行分发、同步以及一致性管理的挑战。

在本文中,我们将简要介绍我们对大模型存储选型的思考、开发中不同阶段的工作负载需求,包括成本、性能和功能等多个维度,并探讨如何通过 JuiceFS 帮助用户优化这些环节。

01 大模型存储选型思考

针对数据量大和数据流转复杂的挑战,通常采用的解决方案是建立统一的数据管理也就是数据湖。这种方法通过统一命名空间,在后端存储不同类型的数据,同时前端提供多样化的数据访问方式,有效减少了复杂数据管道中数据流转的难度和成本。

尽管理想很完美,但现实中如何设计一个统一的技术栈或产品来实现这些目标颇具挑战。从支持前端数据处理框架的角度来看,全面兼容 POSIX 协议的文件系统目前看来是最佳方案。

此外,文件系统还需支持庞大的文件规模,既要能支持数百亿文件的存储,又要保证在此规模下的文件访问性能;这就需要文件系统具备高效的元数据管理能力和多级缓存等高性能实现。

还包括对其他容器平台存储编排的支持、灵活的横向扩展能力以及多云多数据中心的数据同步支持等,这些都是选择数据湖产品时需要考虑的关键因素。

以下表格中列举了市场上一些常见的文件系统产品,并进行了多维度的横向比较。希望能为大家选型时提供参考。

JuiceFS 用户在选型过程中,常询问的产品包括 CephFS、Lustre 和 Alluxio ,我们将重点梳理这 3 个产品的差异。

我们从产品定位角度来比较这几款产品。

在对数据处理框架的支持方面,需要更好的 POSIX 协议兼容性。

  • JuiceFS、CephFS 和 Lustre 都是文件系统,因此它们均提供较全面的 POSIX 支持并确保数据在不同环节的强一致性;
  • Alluxio 主要定位为数据编排和加速工具,只部分支持 POSIX 且不保证数据的强一致性,因此适用的使用场景有所差异。

在支持海量小文件方面,挑战主要体现在文件元数据的管理能力和性能上。

  • CephFS 和 Lustre 利用磁盘存储元数据,并通过内存缓存进行性能加速;文件规模达到一定量级后,元数据的管理和使用成本会明显增加;
  • Alluxio 提供基于 RocksDB 和基于堆内存的两种元数据存储方案,用户可以根据自己的需求选择适合的方案;
  • JuiceFS 使用对象存储作为后端,该存储方式具有扁平命名空间、高可扩展性和低成本等特点,非常适合存储海量文件。社区版的元数据引擎采用开放的插件架构设计,用户可以根据文件规模、性能需求、使用习惯等从多种主流的开源数据库中进行选择来作为元数据存储, 而企业版的则是自研了基于全内存实现的元数据引擎服务,针对文件目录树管理,从元数据访问性能、横向扩展能力几个方面进行定制优化。

在多云数据管理方面,需要对用户透明的文件系统级别复制能力。

  • CephFS 和 Lustre本身不具备文件系统级别复制功能,需要通过存储底座的复制功能或者第三方工具来实现,在数据一致性上难以保障,也会增加额外的管理成本。
  • JuiceFS 和Alluxio 都提供了对用户透明的多云、多 region 的数据分发能力。除此以外,JuiceFS 还提供了在镜像站点文件系统可写的能力;

总结一下 JuiceFS 在数据湖存储中的给用户带来的收益:

  • 统一命名空间提供多协议访问;
  • 完整的 POSIX 协议兼容性;百亿级别文件规模的支持;
  • 数据强一致性;
  • 高并发共享访问能力;
  • 不同负载的性能优化。
  • 多云、多region的数据分发和访问。

02 JuiceFS 在大模型场景的实践

环节1:数据集加载

大模型训练的数据集加载有两个关键点:数据集需要被反复多轮遍历,即多个 epoch,并且在每个 epoch 开始之前数据集需要被随机打乱。这样做是为了确保模型能够充分学习数据中的模式,并避免过拟合。接着这些被随机打散的数据集按批次加载到 GPU 显存中进行处理。

由于这些数据处理主要依赖 CPU 资源,因此 GPU 在这一阶段往往处于空闲。用户需要提升数据集加载的效率,如使用缓存,从而减少 GPU 的闲置时间。数据集通常以结构化大文件(如 FFRecord、TFRecord、Arrow)或海量小文件,如未打包的文本、图片、音频等形式。

数据集的随机打散过程需要随机读取数据文件,结合数据集文件的格式,数据集加载流程对于存储的 I/O 模型是大、小文件随机读。随机读 I/O 需要文件系统具备高 IOPS 和低 I/O 延迟,尤其对于海量小文件随机读的 I/O 处理,对文件系统元数据引擎的性能提出了更高要求。另外数据集在一次训练过程被反复读取,那数据如果都能被缓存在高性能存储介质中,实现高缓存命中率,就可以最大程度的提升数据读取的性能。

JuiceFS 的设计目标是帮助用户在性能和成本之间找到最佳平衡。它通过利用对象存储作为后端数据持久化手段,并结合多级缓存加速架构,为数据集加载提供了高 IOPS 和低 IO 延迟,同时保持成本效益。下面列出了使用 JuiceFS 进行小 IO 随机读操作的一些性能指标:

  • 命中本地缓存时,延迟在 0.05 到 0.1 毫秒之间;
  • 命中分布式缓存(企业版特性)时,延迟在 0.1 到 0.2 毫秒之间;
  • 直接从对象存储读取时,延迟超过 30 到 50毫秒。

使用 libaio 可以获得更高的 IOPS,但这需要集成特定框架扩展以支持 libaio 接口。另外,JuiceFS 企业版提供自研的全内存实现的元数据引擎,可以在管理海量文件规模的情况下,保证元数据请求的平均处理时间维持在 100 微秒量级,在海量小文件的数据集加载场景中,也得到很好的性能。另外通过元数据引擎多分区的架构实现,提供了根据文件规模动态线性的扩展能力。

下面附上一组用户的大文件随机读测试数据,帮助大家理解 JuiceFS 的性能。测试数据为单个 1.8TB 文件。

环节2:训练环节 checkpoint 保存

对训练环节的 checkpoint 文件进行保存,是为了在训练中断时可以从最近的 checkpoint 进行训练恢复,避免从头开始,从而节约时间和计算资源。如果采用同步写 checkpoint,那么在保存 checkpoint 的时间窗口内,GPU 会全部等待。通常,checkpoint 写入都是大文件的顺序写。要尽量缩短这个过程的时间,就需要文件系统提供高性能的写吞吐能力

JuiceFS 通过对接对象存储作为数据持久化层,其吞吐上限受限于对象存储的写带宽、专线带宽以及 JuiceFS 客户端所在节点的网卡带宽。JuiceFS 采用分块存储设计,通过增加对对象存储的并发,可以充分利用对象存储带宽,来提升大文件顺序写的吞吐能力。

此外,JuiceFS 还提供写缓存功能。若后端对象存储存在性能瓶颈,可考虑开启 writeback 写缓存,利用 GPU 节点的 NVMe 高速存储介质作为本地缓存,首先将 checkpoint 文件写入本地,然后再异步上传到对象存储,这有助于减少写入延迟。这里需要注意,开启 writeback 写缓存,可能会在一些场景中导致 checkpoint 文件不一致,无法正常加载恢复训练,针对这种情况就需要加载上一个窗口的 checkpoint 文件进行恢复,从而重跑一些训练 epcoh。

根据用户实际应用的验证,在通过 torch.save 向 JuiceFS 文件系统保存 checkpoint 的场景,每个 GPU 单卡启动的一个写 checkpoint 进程, 可以达到 1.7GB/s 或更高的写吞吐,能够满足该场景的性能要求

环节3:训练和推理环节的模型加载

典型的模型加载场景包括:

  • 训练恢复加载 Checkpoint 文件场景:在训练恢复阶段,采用并行训练的节点仅加载其 rank 的分片 checkpoint 文件,所有节点需要加载的总文件大小就是所有分片 checkpoint 大小的总和,加载 checkpoint 文件的效率决定了 GPU 等待训练恢复的时长,直接影响 GPU 的利用率。

  • 推理服务加载模型文件场景: 将训练好的模型部署到推理服务中,一般需要在推理节点启动时,加载完整的模型文件。当推理节点的数量很多时,如运行上千个推理节点,同时加载模型,每个节点都需要从文件存储上去读取一个完整的模型文件,这时会产生非常大的读吞吐,如果网络吞吐成为瓶颈,模型的加载效率就会很慢。如果文件存储是采用公有云上的对象存储,那对象存储的带宽、专线带宽等都很容易成为瓶颈,而且每个节点都从对象存储 get 完整的模型文件,也会产生大量额外的对象存储带宽和调用成本。

值得注意的是,在模型加载环节中,模型文件的不同格式,对存储 I/O 模型的需求不一样。模型文件可能会采用 torch.save 保存的 .pt、.bin 文件,或使用 safetensors 库保存的 .safetensors 文件。其中 safetensors 文件较为特殊,加载时会采用 mmap 方式,对存储的 I/O 模型为大文件的随机读;而其他格式文件的加载则是大文件的顺序读。针对这两类不同格式的模型加载的需求,来看一下 JuiceFS 的方案。

JuiceFS 对 safetensors 格式模型文件加载(大文件随机读)的优化实践:包括提前将这些文件预热到缓存中,以获得高文件读取 IOPS 和低延迟的性能。最终性能将受限于缓存介质的 IOPS 能力和JuiceFS 客户端的网络环境。此外,为减少预读可能导致的读放大问题,考虑关闭 prefetch 预取功能。若需进一步优化 safetensors 文件的加载性能,可以在模型加载前预读文件到内核的 pagecache 中,这一步骤能显著提高 safetensors 文件的加载速度,性能相较于从 JuiceFS 缓存中读取可提高一个数量级。

JuiceFS 对其他格式模型文件加载(大文件顺序读)的优化策略 :JuiceFS 通过预读功能来提升大文件顺序读的性能。此功能可以预测并提前加载未来可能请求的数据块到内存中,通过合理配置预读窗口大小并增加访问对象存储的并发数,充分利用对象存储带宽。最终读吞吐性能的上限主要受限于对象存储的带宽或到对象存储专线的带宽。此外,预热 checkpoint 或模型文件到缓存中可以进一步提高性能,例如 8 卡单机加载 checkpoint 的情况下,在网卡带宽足够的情况下,热读的吞吐量可达到 10GB/s 以上,足以应对用户对该场景的性能要求。另外在推理加载模型文件的场景中,通过预热的方式,只需要缓存集群从对象存储上 get一次完整的模型文件,后面推理节点从缓存中读取,最大程度提升了读吞吐能力,也节省使用公有云对象存储的带宽和调用成本。

环节4:混合云架构数据快速分发的需求

随着大模型的普及,GPU 算力成为一种稀缺资源,尤其是在国内环境中,不像过去 CPU 资源那样容易获取。尤其是需要进行通用模型预训练的用户,常常需要处理跨地域和多数据中心算力协作的问题,这就要求存储资源能够根据算力的地理位置灵活配置,确保数据能随算力分布。

JuiceFS 通过其镜像文件系统功能,实现数据从主站到其他镜像站点的自动同步,元数据同步在确定的时间窗口内完成,同城可以达到毫秒级延迟,跨城市在国内约为 10-30 毫秒,在指定时间窗口内,镜像站点保证数据一致性,实际数据复制在后台异步完成。一旦元数据同步完成,数据即可在各站点访问;如果实际数据尚未完全同步到镜像站点,JuiceFS 客户端会自动将请求路由到主站点的对象存储以获取数据,虽有性能损失,但数据访问不会中断,并对用户透明。

在这种混合云架构中,为了更好地控制成本,建议在一个地点部署对象存储,各站点本地部署足以容纳训练集数据的缓存集群,以实现近100% 的缓存命中率。数据预先从共用对象存储预热到各站点的缓存集群中,这种部署方式是AI场景中用户最广泛采用的。

下图是一个用户的一主三镜像站点的混合云算力集群架构。每个站点都通过缓存组实现了数据访问的加速。

值得一提的是,在最新企业版中,JuiceFS 新增了镜像集群的可写功能。在多站点训练场景中,当各个镜像站点的训练集群分别保存 checkpoint 文件并写入到各自站点的 JuiceFS 文件系统时,将自动写到主站点文件系统的元数据和对象存储,并进行其他镜像站点的分发,这里元数据的复制是同步的,保证在镜像集群写数据的版本一致性。这一功能简化并加快后续训练任务的 checkpoint 恢复以及推理服务模型加载的流程,更好的保证在镜像集群训练和推理的数据一致性。镜像站点的写入性能此时会受到站点间专线网络状况的影响

03 小结

针对大数据量和复杂数据流转所带来的挑战,常见的解决方案包括建立统一的数据管理系统,即数据湖。这种策略通过采用统一的命名空间来在后端存储不同类型的数据,同时在前端提供多样化的数据访问方式,从而有效降低了复杂数据管道中数据流转的难度和成本。

在文件存储系统的选择上,用户可以从基础架构、性能、兼容性、以及易用性等多方面考虑,以选定最适合的产品。JuiceFS 利用廉价的对象存储作为后端数据持久化手段,并结合多级缓存加速架构,为用户提供提供高成本效益的存储方案。

此外,文章还探讨了在下述关键环节, JuiceFS 实践与优化策略。

  • 数据集加载:数据集大文件和海量小文件的随机读,需要需要文件系统具备高 IOPS 和低 I/O 延迟,以及高性能的元数据访问;
  • 训练环节 checkpoint 保存:checkpoint 大文件的顺序写,需要文件系统提供高性能的写吞吐能力;
  • 训练和推理环节模型加载:safetensors 格式 checkpoint/ 模型大文件的随机读,需要需要文件系统具备高 IOPS 和低 I/O 延迟;其他格式 checkpoint 模型大文件的顺序读,需要文件系统提供高性能的读吞吐能力;大规模推理节点并发加载模型的带宽瓶颈问题;
  • 混合云架构数据分发:可保证数据一致性的镜像读写功能,满足多区域协同训练中的数据集分发、checkpoint 保存,多区域推理服务的模型部署。

希望这篇内容能够对你有一些帮助,如果有其他疑问欢迎加入 JuiceFS 社区与大家共同交流。