LLM训练的存储需求：训练数据和检查点

发表于： 分类于：
字数：88 阅读：≈ 1分钟 浏览：

LLM训练的存储需求主要涉及两个方面：

  • 训练数据：用于更新模型权重和促进模型收敛。
  • 检查点：将模型权重从GPU内存保存至持久化存储。

训练数据

Transformer模型训练本质上是计算密集型的，这意味着所需的数据量相对较小。经过分词处理后的二进制训练数据，每个英语tokens约占4字节，每万亿tokens约需几TB存储空间。

当前领先的模型通常使用数万亿tokens进行训练，相当于数十TB的文本数据。

检查点

LLM检查点的存储需求与模型规模相关，而非训练集群大小。可以通过假设每个参数占用16字节来粗略估算检查点大小。

同样，检查点的性能需求也主要取决于模型规模。通过采用异步多级检查点策略，可进一步降低性能要求。

在大规模系统中，异步检查点机制被广泛采用。在此方案下，GPU仅在将检查点数据从GPU内存复制到主机内存时被阻塞。随后GPU继续计算，而CPU异步将数据刷新至非易失性存储。字节跳动的MegaScale系统采用了这种方法。

利用多级非易失性存储并在不同时间间隔进行检查点，可进一步减少带宽需求并提高可靠性。微软的Nebula框架采用以下策略：

  1. 同步将检查点保存至CPU内存
  2. 异步复制到相邻节点的本地SSD，防止单节点故障
  3. 异步复制到对象存储，防止多节点故障、支持回滚及长期存储

在实际应用中，字节跳动使用HDFS进行检查点存储。

快速检查点与恢复（字节跳动MegaScale）

在识别并剔除故障机器后，驱动程序需要通过加载最近检查点的模型权重和优化器状态来恢复训练。确保最新检查点尽可能接近故障发生时的训练进度状态至关重要，以最小化计算和时间损失。这要求我们提高训练期间的检查点频率。然而，我们还需要减少检查点过程引入的延迟，特别是阻碍训练进度的关键路径时间，因为这会影响整体系统吞吐量。

为实现快速检查点，我们引入了一种优化的两阶段方法。第一阶段，每个 GPU 工作节点将其片上状态写入主机内存，然后继续训练过程。通过优化PyTorch的序列化机制并使用固定内存，借助高速PCIe带宽，该过程可缩短至几秒钟，从而最小化对持续训练过程的干扰。第二阶段，后台进程接管，异步地将状态从主机内存传输到分布式文件系统（在我们的部署中为HDFS）进行集中维护。这种将操作解耦为两个阶段的方法使GPU工作节点能在转储状态后几乎立即恢复训练，而更耗时的HDFS写入过程则被卸载到一个单独的非阻塞进程中。

在从检查点恢复的情境下，由于没有最新检查点就无法开始训练，因此这一过程处于关键路径上。瓶颈在于HDFS的带宽，特别是当每个GPU工作节点需要读取其对应的状态分区时。为缓解这一瓶颈，我们提出了一种优化的数据检索策略。我们认识到多个GPU工作节点通常共享相同的状态分区，例如同一数据并行组中的工作节点。因此，我们指定组内的单个工作节点从HDFS读取共享状态分区，从而线性减少负载。该工作节点随后将状态分区广播给所有共享相同数据的其他GPU工作节点。这种方法有效缓解了HDFS的带宽限制，显著缩短了恢复时间。

LLM中的检查点问题

随着深度学习领域的快速发展，数据集和模型规模呈现爆炸性增长。当前数据集规模可达PB级，单个模型参数量也可达数百GB。这使得模型无法完全加载到标准GPU显存中，凸显了高效并行化和可恢复性的重要性。

全维度并行化

对于LLM等大规模深度学习模型，数据和模型规模已超出单个设备的内存容量。以GPT-3为例，其模型大小超过500GB，而典型GPU仅有80GB显存。且单GPU训练GPT-3需耗时300余年。因此,多维度并行化成为训练和微调大模型的必然选择。

基于斯坦福大学、NVIDIA和微软研究院的研究，结合三种并行化方式可显著提升LLM训练效率：

  • 数据并行：在多个计算设备上复制整个模型，并分配训练数据。实现简单但内存消耗大。
  • 模型并行：将模型划分为离散层或张量，分布到多个设备。实现复杂但更节省内存。
  • 管道并行：将训练过程拆分为小步骤，在不同设备上执行。可能增加延迟,但能提高吞吐量。

LLM

检查点与可恢复性

即便采用并行化，完成一次模型训练仍可能耗时数月。因此，定期执行检查点操作以保证系统状态可恢复性至关重要。检查点通常在每个训练周期后或固定训练步数后执行。这不仅可以在训练中断时恢复进度，还允许回溯调整超参数，确保模型训练的可重复性。

在AI架构设计中，必须充分考虑检查点操作的支持。尽管AI模型训练主要受GPU限制，检查点过程却是I/O密集型操作。写入是检查点的主要负载，而读取则是恢复过程的瓶颈。因此，AI架构的I/O性能直接影响检查点和恢复效率。

检查点大小：深入分析

近年来，随着AI技术的快速发展，基础设施规模成为热点话题。以GPT-3为例，我们可以探讨如何为LLM配置适当的基础设施。

以下是一个具体配置示例：GPT-3拥有1750亿参数，假设部署于1024个GPU(128台NVIDIA DGX-A100，每台8个GPU)。采用三种并行化策略：

  • 张量并行：将500GB模型分散到8个GPU上。
  • 管道并行：在8个DGX节点(八元组)上复制模型。
  • 数据并行：16组8个管道并行系统(共128节点)。

每GPU配置1线程，每DGX 8线程，每八元组64线程。这种配置支持高效的检查点写入。

注意：常见误解是集群中每个GPU都需要进行检查点操作。实际上，只有一个八元组需要检查点。

该配置在Megatron模式下可实现线性扩展，适用于最大1万亿参数的模型，FLOP效率接近理论值的50%。

结论：检查点大小仅取决于模型大小，与数据量或GPU数量无关。配置LLM基础设施时，关键是了解模型大小并确保足够的I/O带宽。

检查点数学分析

下表展示了三个主流LLM的模型大小与检查点状态大小的线性关系：

LLM

检查点状态大小仅依赖模型大小，与数据集大小、检查点频率或GPU数量无关。

检查点操作详细分析

以3072个GPU(384台DGX)训练1万亿参数模型为例：

LLM

读取检查点：受限于存储系统1 TB/s读取带宽

写入13.8 TB检查点：273 GB/s，仅达存储系统写入带宽40%

每DGX约8 GB/s，受8线程限制

写入时间是关键考虑因素。示例中512个GPU需0.53 GB/s/GPU带宽，即4.3 GB/s/DGX。GPT-3检查点时间约50秒，每4小时一次，开销0.3%，远低于5%建议值。

恢复操作

所需读取带宽 = 数据并行度 × 写入带宽 = 6 × 273 GB/s = 1.638 TB/s

写入带宽为恢复所需读取带宽的17%

存储系统需提供1.64 TB/s读取和280 GB/s写入带宽，才能实现50秒的最优恢复时间(13.8 TB模型)。

检查点误区澄清

市场上存在每GPU需1 GB/s带宽用于全集群检查点的误解。实际上，仅管道并行集中的GPU需要此带宽。例如，1万亿参数模型的273 GB/s写入带宽仅对应0.53 GB/s/GPU(512个GPU并行)。

检查点只需记录一致的系统状态，不依赖数据，而依赖模型权重和偏差。

结论

构建 AI系统 时，应依靠严谨的数学和科学分析来驱动决策。考虑到系统规模，微小错误可能代价高昂。务必依赖准确的技术知识和计算，避免粗略估算或带有偏见的指导。