超大规模 AI 模型训练使用对象存储

HPC 专家 Glenn Lockwood 表示，最大规模的 AI 语言模型训练正在使用对象存储，而非文件存储。

AI 模型训练需要非结构化数据。大多数非结构化数据都存储在文件中。因此，大语言模型 (LLM) 训练需要访问文件数据，并需要高性能计算 (HPC) 中使用的并行文件系统，以比串行文件系统更快地向处理器传递文件数据。但是 Lockwood 表示这种说法并不准确，因为 LLM 模型训练各阶段的特点实际上更适合使用对象存储而非并行文件系统。

Lockwood 是 Microsoft 的 AI 基础架构架构师，曾在全球最大的超级计算机之一工作。他写道："在很多人的认知中，超级计算机和并行文件系统就像豌豆和胡萝卜一样密不可分，以至于认为不使用大规模并行文件系统就能运行大规模并行计算任务这种想法显得如此违反直觉，以至于难以置信。"

他提到了 LLM 生产的四个阶段：

数据摄取，"爬虫从互联网抓取原始 HTML、图像、视频和其他媒体。这些原始数据被索引并存入数据仓库。对于前沿模型来说，规模可能达到数百或数千 PB 的数据。"

数据准备，"将原始数据转换为标记化数据。这本质上是一个巨大的数据分析问题，使用 Apache Spark 等框架通过文本和图像处理管道对互联网上的原始垃圾数据进行过滤、去重和清理。数百 PB 的输入数据会被减少 10 到 1000 倍。"

模型训练，"将标记化的数据分批输入到运行在大型 GPU 集群上的 LLM 中。在处理数据的过程中，模型权重会被更新，这些权重会被检查点保存到存储中。"

模型部署和推理，"最终模型被复制到大规模推理服务器集群中，前端的 Web 服务将 REST API 请求转换为在 GPU 上运行的实际推理查询。"

在他看来，每个阶段的 I/O 模式都不需要并行文件系统。数据摄取"只是从互联网抓取 HTML、图像或视频流并将它们打包到数据容器中。在将网页打包到这些文件时，会同时建立一个单独的索引来存储网页的元数据和位置信息。数千个虚拟机可能完全独立地执行这些任务。"这是一次性写入数据，适合对象存储的不可变性。

要点在于"虽然可以将每个抓取的 HTML 页面存储在并行文件系统的文件中，但访问这些文件会非常慢 - 对所有数据进行完整爬取需要扫描数千亿个小文件。" Lockwood 认为"最好在对象存储之上实现数据容器，并使用分布式键值存储来存储索引。"

数据准备涉及"运行类似 Apache Spark 的管道，以完全并行的方式处理所有原始数据...每个处理任务可能一次性从对象中读取几百 MB 的数据，在内存中处理，然后一次性写回对象。文件系统在这里没有优势，因为每个任务只读取一次和写入一次，而不是在单个对象内来回跳转。"

数据去重"需要将每一份数据与所有其他数据进行比较。"这个 I/O 密集型步骤"通常在靠近摄取数据对象存储的集中位置进行，使用基于 CPU 的完全独立的超级计算机，在 GPU 训练开始之前就完成。"

Lockwood 断言"购买便宜的对象存储和 CPU 集群比购买昂贵的文件系统并浪费 GPU 节点进行简单的并行文本处理任务更具成本效益。"

他说，在 HPC 系统中，"快速检查点和重启的需求是创建并行文件系统的主要驱动力。"

AI 模型训练阶段也需要快速检查点和重启。然而，尽管"并行文件系统当然可以用于训练，但它们并不是在数万个 GPU 上进行训练最具成本效益或可扩展性的方式。"

LLM 训练涉及大量重复："无论是在一个还是一千个节点上的 GPU 训练模型，都遵循一个简单的循环 (在 LLM 训练术语中称为'步骤')，这个循环会不断重复。"

1. 将一批标记化数据加载到 GPU 内存中 2. 数据通过神经网络处理，调整模型权重 3. 所有 GPU 同步其更新的权重

第一步的 I/O 负载与传统 HPC 作业不同，首先，"数百万个小型文本或图像文件在 GPU 处理之前就被打包成大型对象。"其次，标记化数据的实际量相当小。一个有 4050 亿参数的标记化 Llama-3 模型"使用 16,000 个 GPU 训练，在整个 54 天的运行过程中只需要从存储中加载 60 TB 的 token。平均每个 GPU 在整个过程中只处理 3.75 GB 的 token。"

这并不是一个很大的 I/O 挑战。此外，标记化数据可以放入 GPU 的本地闪存中，Hammerspace 称之为 0 级存储。Lockwood 指出"使用节点本地 NVMe 可以使存储容量和存储性能与 GPU 性能线性扩展。"

他说："由于两个 GPU 永远不需要读取相同的输入 token，因此在训练循环内部永远不需要在节点之间复制输入 token。"

"在训练过程中加载输入 token 不需要超高带宽或超高容量的并行文件系统。"

检查点写入的实际 I/O 负担也比我们想象的要小。"与科学 HPC 作业不同，检查点大小不会随作业规模而扩展；405 亿参数模型在 16,000 个节点上训练的检查点大小与在三个节点上训练的检查点大小相同。这是因为每个训练步骤后都会进行全局同步，使每个数据并行副本的模型保持相同。只需要保存一份模型权重，对于最先进的 LLM 来说，这个数量不到几百 TB。"

他引用了 VAST Data 的数据说："即使是万亿参数模型，在 3,072 个 GPU 上训练时，使用一个适中的 273 GB/s 文件系统也能达到 99.7% 的前进进度 (只有 0.3% 的时间用于检查点)。要达到这种性能水平不需要并行文件系统；例如，基于 HDD 的 Azure Blob 在使用 IOR 进行大规模写入基准测试时，实现了超过 1 TB/s 的性能。"

做法是将检查点写入 GPU 本地存储以持久化数据，然后异步复制到邻近 GPU 的本地存储以便快速恢复，并迁移到共享存储以进行长期保留。这涉及大块写入，Lockwood 总结道："这种适度的写入带宽和简单的顺序大块写入非常适合对象存储。这并不是说并行文件系统在这里不能工作，但这种检查点方案并不需要目录结构、细粒度一致性语义或其他会增加并行文件系统成本的复杂功能。"

模型部署和推理阶段的 I/O 要求也相当简单：

1. 为推理配置 GPU 节点时，需要尽快从共享存储加载模型权重。 2. 使用 LLM 搜索文档时，需要向量数据库来执行相似性搜索，用相关文档增强 LLM 查询。这是 RAG 的基础。 3. 通常使用键值缓存来减少推理管道不同部分的延迟，存储上下文，包括对话或经常访问的上下文文档。 4. 随着推理需求的变化，不同的模型和权重可能在单个 GPU 服务器中进行交换。

"并行文件系统对这些需求都不特别有用；只有在加载和重新加载模型权重时 (1 和 4)，它们的高带宽才有益处。但就像分层检查点一样，这些 I/O 操作是全对象、只读复制，非常适合对象 API。这里不需要复杂的目录结构和强一致性。"

"这四个步骤中的每一个的 I/O 模式都很适合对象存储，因为它们主要是一次性写入和整文件事务。当然可以使用并行文件系统，工作负载也会受益于它们提供的高带宽。但是它们带来了不必要的功能成本 - 无论是实际成本 (设备或专有软件) 还是隐含成本 (分配人员来管理并行文件系统调试的复杂性)。"

"并行文件系统并不差，它们也不会消失。但是它们对于训练前沿模型来说并非必需，而且...地球上一些最大的超级计算机就是设计成不需要它们的。"