大模型里的Checkpoint和存储方案二三事

大模型通常使用分布式架构搭建系统来进行训练，而分布式架构在设计时就已经把意想不到的各种故障当做寻常事件，并且要建立有效应对机制， Checkpoint 断点也是其中一种。

checkpoint 机制我认为可以类比为快照，它定期保存某些个时间点的完整模型状态的数据，形成如果发生故障，我们就可以调取最近的快照数据，将模型重新恢复到上个检查点的状态，恢复训练。
和快照类似， checkpoint 设置得越多越密集，发生故障后的损失越小；但反过来，太多的快照对系统也有影响和要求，比如多“快照”的性能要求、对系统本身正常工作的打扰，快照数据存储空间都是要考虑的因素。

刚开始大模型的存储方案是五花八门的，放本地盘或是弄一台传统 NAS 都有人试过，这样做的时候，在数据预处理和第一次数据载入的环节，即使慢一点也可以忍受，但在模型规模和训练数据量增长后就不行了。规模变大导致训练周期变长，故障风险升高，因此 checkpoint 设置从无到有到很多，频繁的“打快照”（每次需要写入的 checkpoint 数据从几十 GB 到几个 TB ），导致此处成为对存储要求最高的环节。

针对这样的应用特征进行分析，对存储的需求可以总结为主要是超高性能带宽和高并发，此外存储本身的高可靠和高可用也是必须的。

这个时候，使用 NFS 协议的传统 NAS 存储性能就远远不够了。业界在 2023 年上半年时，还在互相交流在训练节点（服务器）本地多放几块 NVme 的 SSD 盘，突然就获得了性能上的大幅提升，大家惊喜的发现，做 checkpoint 不用再等个半天一天了，用我尊敬的某位老院士的原话：“嘿！一会就得！”。

但是还有个隐患未解，那就是本地盘的可靠性无法与存储系统相提并论。所以到了现在，使用 NVMe 协议的分布式全闪已经是最受认可的解决办法，同时，从存储种类来看，块存储，文件存储和对象里，目前更适配的是文件存储，对象太慢，块没机会。

国外一家科技媒体明确提出：“ We think this will primarily benefit file access all-flash arrays using NVMe protocols and with GPUDirect support for Nvidia GPU servers. Suppliers suchas Dell, DDN, NetApp, Pure Storage (AIRI) and VAST Data are well-positioned to capitalize on this. Panasas sees an opportunity in edge AI. Object storage and block storage suppliers are not so well-positioned.
”
除了硬件，并行文件系统软件的产品和一些数据管理软件如数据仓库也有不同的厂商在尝试，例如 IBM 的 GPFS （现在改名叫 storage scale ），据说某大厂在优化了以后用得还是满意的，不知道具体情况了。
除了找已有的产品搭建解决方案，业界仍然在持续努力中，例如微软针对 checkpoint 的应用场景，提供了一个新产品工具Nebula，号称能把checkpoint的时间从几小时降到几秒。 “Nebula is a fast, simple, disk-less, model-aware checkpoint tool in Azure Container for PyTorch (ACPT).“
链接如下：
https://learn.microsoft.com/en-us/azure/machine-learning/reference-checkpoint-performance-for-large-models?view=azureml-api-2

学术界也有进展： 2024 年 1 月，来自爱丁堡大学和新加波南洋理工大的一篇联合论文里，提出一种 serverless LLM 方法，利用一种加载优化的格式设计，结合高效的多层的 checkpoint 加载系统，能够实现更快速的 checkpoint 加载 ;
论文链接： https://arxiv.org/abs/2401.14351https://arxiv.org/pdf/2401.14351.pdf

相关技术还在持续更新， 今天就先写这两三点