AI的存储需求：训练和Checkpoint

AI的存储需求：训练和Checkpoint

随着AI技术的快速发展，对GPU的需求日益增长，这推动了对高性能存储系统的需求。本文将探讨AI模型训练和Checkpoint过程中的存储需求，包括训练数据读取带宽、Checkpoint写入带宽以及恢复读取带宽等关键指标，并通过具体案例进行分析。

AI的繁荣正在推动对GPU的惊人需求，这导致需要最大限度地提高它们的利用率。

最大化GPU利用率需要平衡计算、网络和存储性能。上图中的“东-西向”网络是GPU服务器节电之间的互连；而“南-北向”网络是连接到存储，用来读取训练数据，写入模型相关的内容。

存储在整个人工智能生命周期中扮演着重要角色。在数据准备——训练和调优——推理这3个阶段，本次讨论聚焦在第二阶段。

AI模型训练使用训练数据批量更新权重，并设置定期检查点（checkpoint）以进行恢复。

关于训练读取性能需求的信息有限；GPU基准测试（具体是MLCommons的MLPerf Training benchmark）提供了一种估算的方法。

训练数据存储的读取带宽需求差异很大；这取决于模型的计算约束和输入大小。

在上表列出的应用类型中，GPT3 LLM大语言模型的175B参数量、305GB的C4数据集都是最大的，但它需要的训练数据读带宽却很小（可能是平均值？）。原因应该是他的数据对象输入大小只有8K字节，使用32个H100 GPU的模型吞吐量（每秒训练样本数）只有19.5K，在这一环节中可以看作是偏计算密集型。

相比之下，3D U-Net三维图像分割的模型参数量只有19M，训练数据集也只有27GB，但输入大小（每个图像）达到92MB，并且每秒训练样本数为463，所以训练数据读带宽需求高达41.6GB/s。

另一种比较典型的Resnet-50图像分类，模型参数量也不大，输入大小中等，但每秒训练样本数较高达到55K，所以训练数据读带宽也达到41.6GB/s。

示例：由DLIO基准测试模拟的ResNet-50训练。

通过NFS（文件存储）读取训练数据会产生连续的顺序IO流，IO大小在64KB到256KB之间，且与批量大小无关。

通过S3（对象存储）读取相同的ResNet-50训练数据会产生更大的顺序20-50MB IO操作。

当多个模型实例在同一服务器上训练时，NFS会从操作系统的页缓存（Page Cache）中受益。

检查点（Checkpoints）会定期保存训练状态，以便在发生故障或中断后恢复训练。

Checkpoint的累积写入带宽需求，取决于模型大小和允许的最大时间。

如上面图表，Checkpoint的大小（GB数）为模型参数（Billion数）的14倍。以175B参数模型为例，如果在2小时的训练中，允许有5%的时间（也就是360秒）用于生成Checkpoint，它所需要的存储写带宽就是6.8GB/s,而在这360秒里相当于GPU算力就闲置了。如果Checkpoint只允许占用1%的时间，那就需要34GB/s存储写入性能，以此类推……

从Checkpoint恢复，需要将保存的状态还原到所有的GPU。

上图是不是有个笔误，好像不是5分钟，而是5小时？

Checkpoint的累积读取带宽（BW）取决于模型大小、数据并行性以及允许的最大时间。

还是以175B模型为例，检查点大小为2450GB，在5小时内恢复1个模型实例只需要0.14GB/s读带宽；而若是16个模型实例（数据并行度）就要2.18GB/s的读带宽。为什么把读取的时间放宽到5小时呢？因为Checkpoint需要定期不停地做；但用到恢复的时候相对较少。

GPU集群运行多个工作负载，依赖于对数据的平等访问，并要求可扩展的存储性能和容量。

AI存储需要在整个AI生命周期中执行和扩展。

还要考虑仍然适用于AI 的传统企业级存储要求。包括：数据保护、高可用、压缩和去重、静态加密、多协议数据访问、远程和混合云复制、安全和合规、长期归档存储，以及数据生命周期管理等。