大模型存储实践：性能、成本与多云

2024-04-07

苏锐

大模型应用领域的迅猛发展，也推动着基础技术领域持续探索和进步。文件存储服务在 AI 基础设施中成为不可或缺的重要部分。

在过去 18 个月的时间里，JuiceFS 团队与 MiniMax，阶跃星辰，智谱 AI，面壁智能，零一万物等大模型团队展开了交流与合作，已经支持了多家客户生产环境中数千卡的训练任务。

在这篇文章中，我们将分享大型语言模型在存储领域面临的一些挑战与 JuiceFS 在服务这些场景时的实践经验，为相关企业提供参考。

01 存储系统性能与成本之间的平衡

在大家刚开始投入到预训练模型时，有这样一种观点：GPU 很贵，相比之下存储的成本忽略不计，可以直接选性能最好最贵的存储方案。典型的高性能文件系统有 GPFS、Lustre、Weka，以及其他高性能 NAS 等。这些系统通常依赖全闪存（NVMe）和高性能网络提供极致性能。

同时，当算力、数据与团队投入都增大的时候，我们又看到了几个新的事实：

零一万物在其最新发表的论文《Yi: Open Foundation Models by 01.AI》（以下简称《Yi 论文》），其预训练数据集包含 3T tokens，通过 BPE tokenization 处理，每个 token 大约占 2Bytes，这意味着 3T tokens 大约等于 6TB 数据。然而，准备可用于正式训练的数据集的过程包括数据抓取、清洗、转换等多个前置步骤，涉及大量的实验。这些实验处理的数据量通常是正式训练数据集的 100 倍以上。随着团队规模的扩大，将产生更多实验结果和中间数据，加上各种模型的 checkpoint 和日志数据，预训练环节总数据量预计将达到 10PB 到 100PB。
正式训练环节，如上文推算，虽然数据集规模固定为 6T，企业可以将全部数据存储于高性能存储系统中。但是，高性能文件系统的性能都与容量是关联的。例如，每 TB 容量提供 250MBps 吞吐，也就是说仅仅把 6TB 数据集存在高性能文件系统，仅能提供 1500MB/s 的吞吐，如果要达到训练所需的 I/O 性能，需要扩大高性能文件系统容量。

于是，出于成本考虑用户通常不会将所有数据仅存储于之前提及的高性能文件存储系统中，如阿里云的 CPFS，其按官网报价（1.4 元/GB/月）来算，10PB 数据的月成本将达到千万级别。还需要注意的是在容量规划方面这类方案不能像 S3 那样弹性使用，扩容很不方便。

用户开始采用对象存储（约 0.13 元/GB/月）与高性能文件存储相结合的策略。这种做法虽然成本更低，但随之而来的是需要额外人力和时间去处理两套存储系统间的数据同步、迁移和一致性管理等任务，这些工作不仅过程繁琐，而且与追求高效率的目标相悖。

那 JuiceFS 如何解决上述性能和成本的问题呢？

经济的对象存储：JuiceFS 将对象存储作为数据持久化层，显著降低了存储成本。
弹性容量管理：它提供类似于弹性文件存储的容量规划能力，弹性扩缩容机制增强了效率，免除了手动管理的需求。
高性能：相较于高性能的专用文件系统，分布式文件系统的性能是用户最关心的问题之一。JuiceFS 通过多级缓存加速架构为预训练提供充足的读写吞吐能力。用户生产环境中的 I/O 吞吐量监测数据，峰值超过了 340GB/s。更多 JuieFS 性能调优策略，可参考：千卡利用率超 98%，详解 JuiceFS 在权威 AI测试中的实现策略。

成本方面，JuiceFS 企业版相较于高性能文件存储 AWS FSx for Lustre，价格仅为其 20%；而缓存层所需要的硬件资源可使用 GPU 节点配备的 NVMe 驱动器，没有产生额外成本。由此，用户能够以仅 20% 的成本获得一个具有弹性的高性能文件存储解决方案，其性能与传统的高性能文件存储系统相匹敌，且具有更优的扩展性。此外，还免除了用户需考虑数据与对象存储之间的同步、迁移及一致性管理等问题。

02 为什么 POSIX 在 AI 训练中必不可缺?

准备高质量的训练数据是构建出色基础模型的基础。数据准备本身是一个复杂的流程，正如《Yi 论文》中所展示的那样：

每个环节的数据处理需求都是不同的，而且这个过程在今天仍然没有一个统一的范式，数据工程师仍在不断实验。

数据工程师几乎都在用 Python，在并行处理中会用到 Ray，如果使用 Spark 也大多通过 PySpark 编程。这些操作的灵活性和高效性要求底层文件系统具备 POSIX 兼容性，这样可以比较高效地满足各种数据处理的需求。
HDFS 只支持追加写，无法支持需要覆盖写的数据处理方法，比如 Pandas。同时，HDFS 的 Python SDK 也不够成熟。
S3 等对象存储不支持高效的追加或者修改（只能整体覆盖），不支持重命名操作。目录操作的性能会很慢。有成熟的 Python SDK，但使用上仍然没有 POSIX 的方式简单直接。另外，数据处理工作还可能会遇到对象存储的带宽限制，高并发下可能会遇到 API 的 QPS 限制。
使用 S3FS 等方案挂载 S3 等对象存储时，可以支持 POSIX 方式访问，但很多操作的性能会比我们预期的慢很多。比如对一个文件做覆盖写，需要将它下载到本地进行修改，然后再完整上传，和文件系统中的局部覆盖写是完全不同的。对一个目录做重命名也会遇到同样的问题。（请参考另外一篇文章）
使用公有云的 NAS 产品，可以用 pjdfstest 做一下 POSIX 兼容性测试。另一个问题是 NAS 的性能与数据量是线性相关的，所以使用中可能会遇到当前数据量提供的性能不能满足计算需要的问题。

所以，对于数据工程师，一个全功能的 POSIX 文件系统是数据处理、清洗环节的最佳拍档。JuiceFS 完全兼容 POSIX ，同时也支持 HDFS、S3 API 的访问方式，在数据处理、清洗的整个流程中能很好地支持各种各样的计算负载。

关于 POSIX 兼容性在 AI 训练中的重要作用，可参考案例，韩国国民搜索 NAVER：为 AI 平台引入存储方案 JuiceFS。

03 多云架构：数据同步、一致性管理的挑战

无论是训练或是推理的需求，单一数据中心或单一云区域内的 GPU 资源往往无法满足用户的全部需求。特别是对于面向 C 端消费者的应用，为了提供更佳的用户体验，常常需要在多个地理区域进行部署。在这种情况下，用户面临的挑战包括数据集和模型在多区域之间的分发、同步及一致性管理等。

下图是某用户在最初使用多云架构时的数据管理方案示意图。用户需要面对的挑战有：

对象存储 A 到对象存储 B 的数据同步：在处理对象存储间的数据同步时，虽然可以采用定时同步特定前缀的数据或设计对象更新回调以触发同步，这些方法在小规模数据处理时简单有效。然而，当同步作业面对大规模数据时，这些同步方案的复杂性急剧上升。挑战包括如何管理同步任务的并发执行、确保数据同步的可靠性、任务失败后的数据重建、系统的观测性、流量控制以及数据一致性校验等一系列问题。
高性能文件存储与对象存储之间的数据同步：由于高性能文件存储的容量有限，需要人工决策哪些数据是近期必需的，并安排合适的时间将这些数据从对象存储中复制过来。当存储空间不足时，又必须协调众多团队成员共同决定删除哪些数据，以释放空间。这一过程中，每个人都倾向于保留自己的数据，从而避免它们被删除，这使得是否扩容或是进行团队内部协调成为一个复杂的决策问题。而扩容并非仅仅关乎成本，还涉及到额外的运维工作和资源投入，增加了同步工作的复杂度和管理难度。
两边高性能文件系统中的数据同步：当用户的任务在区域 A 完成执行后，其可能被调度至区域 B 执行。这就要求任务 A 所使用的数据集需要在区域 B 内可获取，而且其先前的执行输出和日志也必须同样可访问。
同步管理与一致性保证的挑战：选择强一致性还是最终一致性依赖于业务需求和技术实现的复杂度；如果选择最终一致性，明确其时间窗口的界定也是必要的，以保障系统的整体可靠性和预期行为。
存储系统差异问题：这些系统在产品性能、使用限制以及管理策略等方面往往存在细微差异，这些差异要求用户采用精细化的同步和管理方法来确保数据一致性和系统效率。

手工维护上述这套数据架构，运维负担可想而见，如果不止两个区域，上面所有的问题将更加复杂。能否让上面这些问题自动化？JuiceFS 使用镜像功能满足了这样的需求，目前已经成了大模型训练和推理业务中的必备功能。

在 JuiceFS 数据镜像方案中，所有主站点的数据变动都会自动同步到各个镜像站点。元数据同步在确定的时间窗口内完成，同城可以达到毫秒级延迟，跨城市在国内约为 10-30 毫秒，跨大洲则为亚分钟级。在指定时间窗口中，镜像站点保证数据一致性。

数据同步完成的时间基本上与在网络环境中传输一个 4MB 文件所需的时间相同。如果数据尚未完成同步，从镜像站访问文件时，JuiceFS 会自动处理其中的 I/O 路径，确保文件可以被访问。这种处理机制可能会导致短暂的性能下降，但不会影响程序正确执行。而当用户在镜像站产生新数据并写入 JuiceFS 时，数据会自动回写到数据源站，此时的写入性能取决于网络状况。更多实践案例：知乎如何在多云架构下使用 JuiceFS 镜像功能？

04 模型加载慢，GPU 等待时间久

模型加载就是将模型完整读取载入显存的过程，在训练启动、训练恢复和推理服务部署时都包含模型加载过程。模型通常是一个大文件，随着模型参数的增加，模型文件的大小也从几十 GB 增长到 TB 级别，通常采用 pickle 或 safetensor 等格式。

加载过程需要从存储系统中单线程顺序读取，影响速度的关键因素是单线程顺序读取时的吞吐量，JuiceFS 当前版本加载模型吞吐性能为 1500MB/s。经过为模型加载场景的优化后，读吞吐可以提升至 3GB/s。

另一方面，以 PyTorch 加载 pickle 格式模型的过程为例，在顺序读取模型文件的同时会完成 pickle 数据的反序列化，这个过程也会消耗时间。在我们的测试中，从内存盘加载 Llama 2 7B 全精度模型，pickle 格式，26GB 大小，吞吐性能是 2.2GB/s。因为内存是最快的存储介质，所以我们将其视为极限值。从 JuiceFS 加载同样的模型，吞吐性能为 2.07GB/s，是极限值的 94%。

基于这样的表现，越来越多 AI 用户在推理业务中使用 JuiceFS 存储模型，加速模型加载，省下了可观的 GPU 成本。这里可以参考 BentoML 的优化经验，他们是一家模型部署 SaaS 服务，要为客户管理大量模型，非常关心模型加载性能。