一、大模型存储背景

Storage For AI

背景

随着全球数据量的迅猛增长，预计到2025年将达到 182 ZB。人工智能（AI）应用正逐步渗透到各行各业，预计到 2026 年 AI 应用渗透率将达到30%，包括生产效率提升、产业转型升级、社会治理优化等等。以深度学习和大模型为代表的人工智能技术正加速推动智能化转型，这不仅为基础设施的算力、存储、网络提出了前所未有的高要求，也带来了一系列挑战。

GPU、TPU、NPU 等异构算力正以每年超过 10 倍的速度增长，而无损以太网技术 400GE 已成为行业标配。在这一背景下，存储逐渐成为制约人工智能应用发展的瓶颈之一。因此，“Storage for AI”（为AI系统设计的存储）逐渐成为学术界和工业界的热门研究方向。

2018年，NSF（美国国家科学基金会）的《Data Storage Research Vision 2025》报告首次提出了“Storage for AI”和“AI for Storage”（AI赋能存储）两个概念。其中，“Storage for AI”关注的是如何通过系统化的设计，使存储系统更好地服务于AI工作负载和数据使用；而“AI for Storage”则关注如何应用AI技术和算法，赋予存储系统更高的智能性、自主性和效率。

挑战

2022年11月 ChatGPT 3.5 的问世，工业界的大模型如雨后春笋般出现，各大厂商开始军备竞赛，学术界的论文一边倒向大模型，GPU、TPU、NPU算力持续增长，存储已经明显的成为问题。主要的挑战在如下几个方面：

数据集规模挑战：过去十几年数据集的大小在指数级的增长，从 ImageNet-21K 开始数据集大小从 1.4TB 到 Sora 训练数据集可能的 20PB。从 Huggingface 的数据集列表也能看到，很多模型的数据集都超过 1TB 。

数据集访问模式的挑战：数据集文件大多数为 100K-200K 左右的小文件，每个 epoch 训练任务会对整个数据集进行一次 full shuffling 之后，对整个数据集遍历一次。传统存储系统是针对大文件的连续 IO 来优化，离散的小 IO 导致空间局部性和时间局部性都较差，对缓存系统的淘汰机制也不友好。

参数规模的挑战：以大语言模型为例，模型参数量平均每2年增长240倍，而 GPU 的内存平均每2年仅增长2倍，容量小且增长缓慢，出现了内存墙问题。

计算与存储（带宽）的失衡：显存（HBM）大小20年增长了100倍，存储带宽的PCIe gen3 从32GB/s 增长到 gen7 的 512GB/s，15年只增长了16倍。GPU 的 FLOPS 却以20年60000倍的速率在增长。随着异构算力厂商的角力，FLOPS 不断提高，数据存储和计算模块之间的差距会越来越大。

异构计算和存储之间依然需要 CPU 来协调：

张量数据进入到 GPU 之前，要进经过一个漫长的链路：远程存储 -> 本地网卡 -> PCIe -> System Memory -> GPU Memory。

这里存在三个问题：1. 链路长；2. CPU 浪费；3.南北向 50GB（PCIe gen5的经验值）带宽限制，数据很难高效、快速读取到 GPU 内。

总结以上的挑战，随着数据量的增长，从 GPU 内存转换到 NVMe/SSD 是必经之路，构建/用好存储成为 AI infra 里一个非常 critical 的部分。冯诺伊曼体系结构里存储就是二等公民，在 HPC（High-performance computing）时代是食之无味，在 AI 时代已经变得不可容忍。Storage for AI 是一个新的发展领域，也带来了前所未有的机会。

业界存储方案

为了解决大模型场景下海量数据的访问需求，高性能文件系统成为当前的解决方案。在过去几年中，我们看到各厂商和开源界的主要方案可以分为 “独立文件系统”和 “文件缓存+大容量对象多级组合” 两种方案。

Note：文章只罗列了其中几个常见的系统。

独立文件系统

GPFS

最初是IBM 为美国能源部需求设计的一套并行文件系统，具有高性能，高吞吐，高度兼容 Posix 语义的特点，时至今日活跃在 HPC 和 AI 存储设施领域。下图为 GPFS shared-nothing mode。GPFS 面向高可靠的硬件环境，对硬件要求高。当前部分云厂商的高性能并行文件系统会基于该方案演进。

GFS

Google 三驾马车之一，区别于 GPFS，GFS 以廉价设备存储数据，不追求兼容 Posix API 而是面向大数据系统特点进行设计，很好的支持了 MapReduce，BigTable 等系统，定义了大数据时代全新的文件系统范式。由于 GFS 不开源，社区实现了 HDFS - Hadoop 项目的大数据存储系统。

近端缓存 + 大容量的组合方案

H云存储

在大模型训练场景，H云采用了三层架构，容量层 + 性能层 + 算力层，提供成本 & 性能的综合优势。容量层 OBS 负责冷备数据的持久化存储；SFS Turbo 用于训练热数据的加速，在训练开始前从 OBS 预热数据到 SFS，以及定期把冷数据持久化到 OBS；而 AITurbo 则通过 SDK 的方式在训练节点提供 IO 上的适配和优化，比如在 FO 时通过内存里 Checkpoint 来加速恢复。

Meta

Meta最近在其大模型存储的论文中介绍了 Tectonic-Shift 的缓存方案，通过基于历史访问加权的方式来动态加载和淘汰缓存中的数据，以支持 1:N 的缓存持久化比。

GCP AI Storage

Google Next' 24 大会上介绍了最近 GCP 上 AI 存储的相关方案，可以看到在 Cloud Storage 上针对不同的训练需求提供文件系统‍、SSD 等不同类型的存储服务。其中 Anywhere Cache 是一个新产品，支持在多个 AI 站点之间提供高速的数据同步和缓存服务，减少跨区域和跨网络带来的数据访问延迟，从而加速训练任务的数据读写。

蚂蚁大模型存储场景的需求

随着大模型的快速发展，算力需求和数据量（数据集和模型）也急剧增长。经过一年多的时间，蚂蚁的 AI 算力中心从主站的几个集群逐步建设为部署在多云环境上的多 AIDC 站点。这个过程中，在训练规模、效率、稳定性等方面对存储提出了更高的要求，主要体现在以下几个方面：

多云环境需求

• 支持多云环境部署，能够在私有云和公有云上部署。
  
• 具备良好的硬件适配性，能够在不同的硬件环境下部署，并且同时有较好的读写性能。
  
• 低成本：对比云上各类存储服务，在性能和成本的综合评估上有优势。
  
• 多云数据同步：支持多云环境下 AIDC 站点之间的数据同步，为算力的全局调度提供必要的基础。
  

性能 & 数据规模需求

• IO 性能：
  ➢万卡 MOE 任务的单次 checkpoint 达到了20TB，存储的写入性能会直接影响训练的效率。
  ➢海量小文件随机读取：在十/百亿规模下，支持高效的小文件随机访问。
  ➢单集群支持 TB 级的吞吐。下图是一个集群的流量监控，峰值会达到1TB+/s。
  
  
  
• PB 级数据管理：多模态单次训练的数据集会达到百 TB 级，甚至 PB 级。如何在数据集规模不断扩展的情况下，始终能够保障数据的高效访问是一个挑战。
  

稳定性需求

持续提升AI训练迭代效率的背景下，存储服务需同时满足高可用性要求，以保障分布式训练任务的全链路稳定性。

• 服务稳定性：百亿文件 & PB 级数据规模的集群服务可用性要达到 99.9% 甚至 > 99.95%，同时具备快速的 Failover 能力
  
• IO 链路稳定性：如何保障写入7个9甚至8个9成功率，让万卡任务的有效训练时长达到 95% 以上
  

二、蚂蚁大模型存储方案

整体架构

随着大模型的快速发展，数据量在急剧膨胀同时读写需求也在不断变化，这些都对存储的 IO 性能提出新的挑战。为了解决这些问题，当前业界大部分公司，e.g：Google，Meta，阿里云，H云，etc，都采用了“计算端基于 GPU 内存 + NVMe 构建的分布式缓存层，全闪存分布式文件/缓存系统，对象存储容量层（作为数据湖底座）”这套解决方案。蚂蚁也同样构建了一套全闪分布式文件缓存系统 PCache，用于支撑蚂蚁内部的大模型训练。

全闪文件缓存层

数据集格式与接口

• 数据格式：支持非结构、结构化等多类型的文件存储和读写。多模态和 NLP 场景主要是非结构化数据，搜推则需要结构化的表类型数据
  
• 用户接口：提供多语言 SDK（Python & Java）和 Posix（FUSE）两类型接口。加速层（从网络近端加速到 AI Native 加速）
  
• 网络近端加速：通过采用 GPU 混部的方式，让存储节点可以使用 GPU 机器上的 SSD 资源，在保证低成本的同时因为近端的网络开销减少而提供高性能 IO。除此之外，存储的整体吞吐可以随着训练集群规模的扩大而自然的扩展
  
• AI Native 加速：在数据读取和 checkpoint 阶段提供面向 AI 数据特性的缓存策略和数据预期优化，以及打散 + 本地 + 异步化的 CKPT 写入机制，来覆盖训练场景的高性能 IO 需求，以及最大化的利用存储性能。
  

基础设施层

• 通用数据集成框架：提供多存储和多云之前的数据集 + 模型数据同步能力。
  
• 云原生存储：分布式缓存系统基于 K8SService, PVC 等资源抽象来构建，在故障发生时具备更好的自动恢复能力，以及更方便的支持自动扩容、负载均衡等功能。
  
• 持久化层：适配多种类型的持久化存储，支持数据的自动预热、穿透、补齐等能力。同时和蚂蚁自研的对象存储 HexS 在回收站、TB 级文件持久化等场景提供整体端到端的高性能存储解决方案。
  

图：PCache 整体架构

在接下来的章节里我们会从部署模式开始，到存储服务优化，再到面向 AI 数据场景的优化来说明 PCache 是如何在蚂蚁的环境里支持高效的大模型训练，并同时具备成本、稳定性和性能方面的综合优势。

GPU混部

从独立部署到 GPU 混部

当前大部分存储服务都是以独立集群的方式部署，但是在多云环境里，各个云厂商提供的硬件环境和存储服务都不尽相同。

主要有以下几个问题：

• 依赖云厂商的存储服务：每个厂商的存储能力不一致，会导致不同云环境的训练效果不一致。同时数据迁移的效率也可能会受到影响，特别是建站或者退租时的大规模数据迁移（依赖全量文件列表获取）。
  
• 自建存储服务：采用独立部署的方式，需要考虑各站点硬件的适配性和对应的运维成本。还有就是网络方面，不同云环境的网络配置不一样，有的环境在数据面有独立的高性能网络，有的则没有；其次就是独立部分的服务一般无法利用东西向的网络带宽。因此面对这些问题，我们希望能构建一个大模型的存储服务，在多云环境里具备良好的硬件环境适配性和扩展性，在成本和性能的综合比较上有优势。
  

GPU 混部

当前大部分的 GPU 机器上都具备4块甚至8块的 NVMe SSD 资源。根据该配置特性，PCache 选择采用混部的模式，把存储服务部署到了 GPU 节点上，减少了独立部署的成本，并通过近端部署减少了网络访问延迟；同时除了南北向网络还可以利用东西向的网络扩展读写吞吐，让存储的吞吐能够伴随着 GPU 集群算力的扩展而自然增长。

稳定性问题

虽然通过 GPU 混部，PCache 获得了成本和性能上一定的综合优势，但是因为 GPU 的高故障率（e.g., 日均 0.1% 坏卡率），会导致混部存储节点的不稳定性，这也是很多产品不选择存储和计算混部的主要原因之一。为了解决该问题，PCache 利用了云原生系统在服务发现和容器编排上的优势，来提高故障时的 FO 能力和效率；同时通过在读写链路上的稳定性优化，来保障故障发生时的读写成功率。

云原生存储

• 基于云原生的方式部署服务：POD 管理计算资源，PVC 管理存储资源。
  
• 故障时的自动恢复：通过 POD name、PVC name 的管理，保障容器重启后对外服务地址不变，数据不丢。
  

集群变更稳定性 - 基于容器编排能力

基于 K8s 的容器编排容易可以让 PCache 更加高效的支持各类服务节点变更的自动化。

• 集群升级：节点只读 -> 节点升级 -> 节点恢复
  
• 故障机迁移：老节点只读 -> 新节点创建 -> 数据迁移 -> 老节点下线
  
• 故障机器下线：感知节点下线 -> 自动触发数据预热
  

数据链路稳定性

为了避免节点故障时，可能导致的数据写入失败或丢失，PCache 从客户端到服务端的数据链路做了端到端的稳定性优化。

• 读写链路优化：采用多链路写入的方式，减少突发故障机下线时的影响。同时增加链接和数据传输重试机制，进一步减少故障节点可能带来的数据丢失概率。
  
• 副本自动补齐：当节点下线时，对于未完成持久化的数据，系统会自动做副本的补齐，保障数据的可靠性。
  
• 数据自动从持久化存储里恢复：最坏的情况下，如果缓存数据所在节点异常下线，数据读取时也会自动从持久化里恢复。
  

数据链路

元数据服务

在多模态的训练中，海量小文件的元数据访问效率始终是摆在大规模训练面前的难题。为了解决百亿文件规模下的元数据性能问题，PCache 采用了在训练侧减少元数据规模和在存储侧提高元数据可扩展性的组合方案。

减少元数据规模 - 数据折叠

为了减少元数据规模，PCache 和多模态数据团队一起设计了小文件折叠的方案，通过离线的数据处理把多个小文件折叠成一个大文件。以下是几种常见的数据折叠场景。

• 按文件数量折叠
  

➢按列存的格式把多个文件折叠为一个文件，同时通过 unifile SDK 解析折叠文件，屏蔽用户对折叠文件的感知。

➢该方案大幅减少元数据数量和读取请求量，线上的多模态任务的数据读取性能提高 2～4 倍。

• 多维度折叠
  

除了数量单一维度的折叠外，现在也出现了越来越多的多维折叠需求，e.g., 卫星图片场景下的时空维度。

提高元数据处理能力 - 元数据存储 & 服务分离（进行中）

当前 PCache 主要通过联邦模式来支持百亿规模的文件存储，但是联邦模式在运维和问题排查上较大弊端，尤其在子集群扩容时。为了解决该问题，PCache 开始做元数据服务和存储分离，支持元数据 sharding & 外置高性能元数据存储（UNS）等多种元数据管理策略，来提高存储自身的元数据处理能力；同时通过 Proxy 组件屏蔽用户对不同元数据存储方案的感知。

数据链路性能优化

在万卡万亿参数的 MOE 训练场景中，单次 checkpoint 写入量会达到 20+TB，为了解决该类型场景下 TB 级数据的高性能读写需求（提高训练任务的有效训练时长），PCache 在数据链路上从客户端到服务端做了端到端的写入性能优化。

客户端性能优化

• 用户态 FUSE：在 Checkpoint 读写场景，为了消除用户和内核态的多次切换或拷贝开销，PCache 采用了拦截的方式，直接通过 shm 来加速中大型文件访问效率。
  
• MetaData Cache：文件 + 数据块组合的元数据 cache 策略提高客户端的读取性能，尤其在大量随机读存在的场景。
  
• 从用户文件并行到数据块写入并行：提高用户低并发场景下的性能。
  
• Worker 选择策略（RR + Load Aware）：减少热点带来的性能下降。
  
• 多路 IO 扩展：支持 TCP/IP，RDMA，本地 IO 多链路写入，扩展带宽；同时 RDMA 也能够有效减低延迟和 CPU 开销。
  

服务端性能优化

• Direct IO：直接写入磁盘，不经过 OS cache，在 CheckPoint 场景有较大性能提升。
  
• 数据分布优化 & 动态副本：➢数据块均衡分布。➢根据访问频率 & 吞吐自动调整副本数，减少热点节点。
  
• 读写线程自动配比：根据盘数，存储容器 CPU 资源做读写线程配比，线程池过大带来的争抢会降低性能。
  

万卡场景性能数据

• 单并发测试：
  

下图展示了 RDMA、用户态 FUSE、客户端并行等优化开启后的效果，对比优化前提高了 3~4 倍的吞吐。

• 多并发测试：
  ➢当前单客户端10~16并发的测试场景下，整体吞吐可以达到 20~30GB/s。
  ➢ 在多客户端的测试场景中，表现出了较好的扩展性，整体吞吐能随着客户端数量的增加而线性增长。300台 GPU 机器(单机配置4块 SSD）组成的混部集群峰值可以支撑 3TB/s 的吞吐。
  

AI Native 存储

除了存储服务自身的能力优化外，为了应对大规模训练场景下的海量数据读写，PCache 也在尝试深入训练任务，向上提供面向训练数据和框架特性的优化策略。支持更大规模的训练数据 - 数据动态淘汰 & 加载

读取性能优化

为了进一步提高数据的读取性能，尽量从高速的存储介质中读取数据，PCache 会构建一套基于面向 AI 数据特性的 Cache 策略，比如基于 shuffle 的伪随机特性，可以通过预测的方式，提前把数据加载到客户端。

提高缓存利用率

对于视频相关的多模态场景下，训练数据集可达到 PB 甚至 10PB 级，如果所有数据都要预先加载到缓存中，势必会带来较大的成本压力。因此为了提高缓存利用率，PCache 会基于训练数据的可替代、可预测等特性，在训练过程中实时和动态的加载和淘汰缓存中的数据，以做到用较少的缓存资源支撑大规模数据的多模态训练。

图：基于数据集切分的动态淘汰 & 加载

我们拉取了一个线上 PB 级集群一周时间的存储水位变化情况。能看到文件会根据历史访问权重做自动淘汰，让集群自动保持在60%的安全水位，减少了额外的空间运维成本，也降低了线上因为存储空间问题引发的故障。

图：PB 级集群的数据自动淘汰效果（过去7天）

Megatron Checkpoint 写入的优化

Megatron 是 MOE 场景训练任务的常用框架。该训练框架在 checkpoint 写入阶段时，优化器数据按 DP（Data Parallel）Group 写入，默认每个 DP 组的 rank0 卡负责数据的 gather 和写入存储。

存在的问题

• 当前 DP 组的 rank0 会调度在同一个节点，在 gather 时网络的写入吞吐较高，同时 cpu 的争抢和 shm 的带宽也会导致 gather 的过程变慢。
  
• gather 完成后，也会导致 checkpoint 写出时打满网络带宽，而与此同时大部分节点的网络则处在空闲状态。比如当前在蚂蚁的某个站点，会因为单节点要写出近 700GB 数据导致 checkpoint 保存节点对训练效率产生影响。
  

按 DP 组打散

和外部云存储团队交流后，可以通过提供 SDK patch 原生的 load/save_checkpoint 方法，采用按 DP 打散的方案来减少单节点内的计算资源和网络带宽争抢，以提高 checkpoint 的整体读写效率。

图：基于 DP 组打散的加速方案

存储管理&生态配套

多云数据同步

在当前训练任务需要经常在多个 AIDC 之间做调度和调配，e.g.，A 站点训练任务产出的 checkpoint 需要迁移到 B 站点做评测。因此数据同步的效率也是全局训练效率的关键一环，而多云同步主要就是解决该场景下多算力中心之间的数据同步效率问题。

• 支持多种持久化存储，能够在不同的云环境之间提供数据同步服务。
  
• 结合蚂蚁的训练任务的管理平台 AIStudio 和蚂蚁数据管理平台 ODC 统一管理多云环境的数据集，避免大量重复数据。
  
• 集成高性能分布式数据集成工具，提高数据迁移效率。当前具备日常 15TB/h 的同步吞吐，并能随着带宽的扩大而增长。
  
• 通过同步过程中的数据校验和同步后的离线旁路核对，保障同步数据的正确性。
  

IO链路巡检 & 监控

除了集群监控外，高可用团队和 PCache 团队共建了端到端的巡检服务，能够发现整个 IO 链路上的异常点。同时基于链路的 trace 日志，分析和产出用户侧的读写成功率统计，以及异常 IO 操作和 checkpoint 写入时间等，用于报警和问题排查。

离线数据分析

通过离线分析的方式快速获取集群的快照信息，e.g. 全量文件列表，异常状态文件信息等。在几十亿文件规模的集群做站点迁移时，离线快照能力可以快速获得文件列表并切分，避免因为 list 性能影响数据迁移效率。

存储管理能力

除了周边生态的服务，PCache 也在不断完善面向训练用户的使用场景的存储管理能力，当前主要支持以下几种：

• 数据集只读：挂载时设置只读，避免数据集被异常修改导致训练任务中断。
  
• 回收站：支持数据删除后的恢复能力，同时保留时间可配置。
  
• 配额管理：具备多级目录的配额控制能力，避免集群内用户之间的相互影响。
  

三、规划

大模型的数据量还在每年以几倍甚至数十倍的速度在增长，蚂蚁的存储在接下来几年需要面对百 PB 甚至 EB 级数据规模的高性能读写需求。未来的数据类型和接口也可能会发生变化，文件不再是存储和训练之间的唯一标准接口，更多类型的数据接口需求会出现。与此同时，如何从提高训练效率到提高训练算力也会是未来大模型存储的一个重要方向。