arXiv:2601.16956cs.DCcs.AI2026-01被引 3

提出新型检查点架构,显著提升万亿参数模型训练效率。

DataStates-LLM: Scalable Checkpointing for Transformer Models Using Composable State Providers

  • 通过可组合的状态提供者解耦状态抽象与数据移动
  • 实现非阻塞异步快照,提升检查点吞吐4倍
  • 适合超大规模语言模型训练团队使用

大型Transformer模型(如千亿参数级语言模型)的快速发展,要求在数千个GPU上采用复杂混合并行策略进行训练。检查点技术对训练容错、暂停恢复、异常轨迹分析等至关重要。现有方案将模型状态视为不可知的二进制块,忽视了其内存位置(GPU/主机)、逻辑对象分片数量、数据类型(张量/Python对象)及序列化需求的三重异构性,导致设备间传输阻塞、序列化盲目和存储I/O竞争等问题。本文提出DataStates-LLM,一种新型检查点架构,利用状态提供者将状态抽象与数据移动解耦。借助参数在前向/反向传播中不变的特性,实现懒加载、非阻塞异步快照。通过状态提供者高效合并异构碎片,并行化元数据序列化与张量批量I/O。我们在256块A100-40GB GPU上评估了最大700亿参数的模型。结果表明,相比现有最优方案,该方法检查点吞吐最高提升4倍,端到端训练时间减少最多2.2倍,有效缓解了超大规模模型训练中的序列化与异构瓶颈。

原文摘要 · Abstract (English)

The rapid growth of Large Transformer-based models, specifically Large Language Models (LLMs), now scaling to trillions of parameters, has necessitated training across thousands of GPUs using complex hybrid parallelism strategies (e.g., data, tensor, and pipeline parallelism). Checkpointing this massive, distributed state is critical for a wide range of use cases, such as resilience, suspend-resume, investigating undesirable training trajectories, and explaining model evolution. However, existing checkpointing solutions typically treat model state as opaque binary blobs, ignoring the ``3D heterogeneity'' of the underlying data structures--varying by memory location (GPU vs. Host), number of ``logical'' objects sharded and split across multiple files, data types (tensors vs. Python objects), and their serialization requirements. This results in significant runtime overheads due to blocking device-to-host transfers, data-oblivious serialization, and storage I/O contention. In this paper, we introduce DataStates-LLM, a novel checkpointing architecture that leverages State Providers to decouple state abstraction from data movement. DataStates-LLM exploits the immutability of model parameters during the forward and backward passes to perform ``lazy'', non-blocking asynchronous snapshots. By introducing State Providers, we efficiently coalesce fragmented, heterogeneous shards and overlap the serialization of metadata with bulk tensor I/O. We evaluate DataStates-LLM on models up to 70B parameters on 256 A100-40GB GPUs. Our results demonstrate that DataStates-LLM achieves up to 4$\times$ higher checkpointing throughput and reduces end-to-end training time by up to 2.2$\times$ compared to state-of-the-art solutions, effectively mitigating the serialization and heterogeneity bottlenecks in extreme-scale LLM training.

大模型训练检查点分布式系统异构优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。