大模型Checkpoint恢复基准：AWS存储方案实测与优化指南

如果你的训练任务在AWS上跑了三天，好不容易推进到第2000步，结果一个Spot实例回收通知下来，节点全没了。重新拉起之后，最想做的事情不是骂人，而是赶紧把Checkpoint读回来，继续跑。但等你真的开始做这件事，你会发现一个扎心的事实：模型训练本身可能不算难，难的是把Checkpoint快速、稳定地读回来并恢复训练状态。这个系列是我在AWS上跑大模型训练和微调时沉淀下来的实操记录，第二篇专门聊Checkpoint的读取与恢复基准——从测试设计、脚本写法、参数推导，到S3、EFS、EBS、FSx Lustre几种存储方案的实际表现，以及恢复慢的时候到底该查哪里。适合正在做大模型训练工程化、MLOps，或者准备在AWS上用低成本资源跑训练的人参考。

1. 大模型Checkpoint读取与恢复基准，到底在解决什么问题

1.1 训练中断是常态，恢复能力直接决定训练工程化成败

大模型训练不是一个“start到end”的直线过程。一个70亿参数模型，全量微调还好，预训练或者更大规模模型，动辄要跑几天甚至几周。这么长的时间里，单点故障、节点被回收、存储抖动、OOM，任何一件事都可能把训练打断。AWS上很多人喜欢用Spot实例省钱，但Spot的代价就是随时可能被回收，虽然系统会提前给两分钟通知，但两分钟只够你保存Checkpoint和优雅退出，根本不够完成任何训练阶段。

这时候Checkpoint恢复能力就成了训练工程化的生命线。保存得再频繁，如果读不回来或者读得太慢，一切都是白搭。我在实际项目里见过一种情况：训练任务每5分钟存一次Checkpoint，结果实例需要重新拉起时，光把30GB的Checkpoint从S3拉回本地就花了十几分钟，再加上加载、校验、恢复分布式训练状态，整个RTO（恢复时间目标）接近半小时。比训练中断本身还难受。

Checkpoint恢复的链路看起来简单：从持久化存储读取模型权重、优化器状态、调度器状态、随机数状态，然后加载进显存，继续跑。但这个链路里，读取环节往往是最容易被低估的瓶颈。尤其是分布式训练场景下，Checkpoint可能是几百上千个分片文件，存储类型、网络带宽、并发策略、文件数量，随便一个变量变了，恢复时间可能差一个数量级。这也是为什么一定要做基准测试——不是凭感觉，“感觉S3很快”，而是用数据说话。

1.2 基准测试到底在测什么：三个核心指标和两种典型场景

我给这套基准定义了三个核心指标：读取吞吐（GB/s）、恢复时间目标RTO、恢复成功率。读取吞吐就是单位时间内能从存储里读出多少数据，这个指标最直观，方便横向对比不同存储方案。RTO是从开始读取到训练状态完全恢复、可以继续迭代的总时间，它比单纯的数据读取多几层，包含了文件清单列举、校验、反序列化、显存分配等步骤。恢复成功率容易被忽略，但非常重要——读一遍没成功，重试又失败，这种问题在文件数量大、并发高的时候经常出现。

测试场景我分成了冷读和热读两种。冷读指实例重建或容器重启后，Checkpoint不在本地任何缓存里，必须从持久化存储（比如S3、EFS、FSx）完整拉取。这是最贴近故障恢复的真实场景。热读指文件已经在本地EBS或实例存储上，只需要读入内存并加载到显存，对应的是训练进程崩溃但机器还在的情况。两种场景分开测，因为它们的瓶颈完全不同。冷读拼的是存储服务和网络带宽，热读拼的是内存带宽和反序列化效率。

除了指标和场景，还要固定几个变量：实例规格、软件框架版本、网络路径、文件布局。否则测出来的数据没有可比性，今天和明天的结果可能完全对不上。后面我会细说这些变量怎么控制。

1.3 为什么必须在AWS上实测：经验估算靠不住

很多人会问，既然有网卡带宽、存储读写速率这些理论值，直接估算不就行了？我的答案是：理论值可以参考，但不能作为UAT依据，必须在真实环境里实测。原因有几个。

第一，AWS的共享型存储服务有各自的限额和突发机制。S3对单个前缀的GET请求有速率限制，小文件一多，你会先撞上请求数上限而不是带宽上限。EFS有Burst Credits机制，短时间内可以飙高吞吐，但Credit耗完之后性能会掉到基线水平。EBS的吞吐和IOPS取决于卷类型和卷大小，一个300GB的gp3卷和一个10TB的io2卷，实测性能差距巨大。这些特性都不是一句“没问题”能带过的。

第二，框架的加载方式直接影响瓶颈位置。PyTorch的torch.load会把整个对象读进CPU内存再反序列化，遇到分片Checkpoint还要做跨进程的reshard；DeepSpeed的Zero-3状态更大，动辄几十GB。就算是同一个模型，用不同版本框架保存出来的Checkpoint结构也有差异，恢复性能自然不同。

第三，网络路径的影响。数据从S3走公网Endpoint还是VPC Gateway Endpoint，跨可用区还是同可用区，延迟和带宽的表现都不一样。我们曾经因为没走Gateway Endpoint，从S3拉数据额外多花了两倍时间，排查了半天才意识到是网络路径的问题。

所以，别省这一步。把基准测试当成训练基础设施的一部分，每次改环境、升版本、换存储之后，都跑一遍，用数据替代直觉。

2. 测试环境与基准设计：先搭台子再跑分

2.1 实例、存储与框架版本怎么选

测试环境的搭建要尽量贴近真实训练环境，否则测出来的数据没有参考价值。我这次使用的实例是p4d.24xlarge，8张A100 40GB，网络带宽400Gbps（约合51.2GB/s的理论上限）。不过要注意，这个400Gbps是实例级别的总带宽，实际到单个S3桶、单个EFS文件系统时，不可能吃满，能跑到几GB/s已经算不错。

存储方面，我准备了四条路线：S3（标准存储桶，使用VPC Gateway Endpoint）、EFS（选择弹性吞吐模式，置备了足够的吞吐量）、EBS（gp3，3000 IOPS基线，125MB/s基线吞吐，测试时开到10000 IOPS和500MB/s吞吐）、FSx for Lustre（做成了与S3关联的数据仓库模式）。每个存储方案都单独测试，避免互相干扰。

软件版本锁定非常关键。PyTorch选2.1.2，DeepSpeed选0.14.4，Hugging Face Transformers选4.39.3，Python版本3.10。为什么版本要锁得这么死？因为DeepSpeed的ZeRO-3 Checkpoint格式在0.13到0.14之间有过不兼容的改动，Hugging Face的模型结构序列化方式也可能变化。你不想辛辛苦苦测出的结果，过两个月因为框架升级就全部失效。我的做法是每次测试都记录一份pip freeze到实验日志里。

2.2 模拟Checkpoint文件结构与三种测试规格

真实的大模型Checkpoint不是单个文件。以7B参数模型为例，混合精度训练下，模型权重大约14GB（BF16），Adam优化器状态至少还要两倍，也就是28GB往上，整体一个Checkpoint接近50GB甚至更多。而且分布式训练时，每个GPU rank会各自保存自己分管的那部分权重和优化器状态，再加上tokenizer、config、调度器状态等小文件，整体文件数量很容易上千。

为了做基准测试，我构造了三种规格的Checkpoint样本，都是模拟真实结构生成的：

生成Checkpoint样本时不计时，只用它来测试读取和恢复。生成方式用PyTorch的torch.save和torch.distributed.checkpoint分别做，因为两者的文件组织方式差异很大。分布式Checkpoint还会生成独立的metadata文件，虽然只有几KB，但在加载时扮演核心角色，读不到metadata整个恢复就失败了。

2.3 测试维度：冷读热读、并发策略与数据采集

测试分冷读和热读两条线。冷读的做法是创建一个全新的临时实例，挂载好存储，从零开始读取Checkpoint并加载。为了让每次冷读的起点尽量一致，我会在测试前卸载存储、清空本地缓存，必要时重启容器。热读则是把Checkpoint预先拉到实例本地，再模拟训练进程崩溃后重新加载的状态。

并发策略是另一个需要控制的变量。从S3拉取数据时，我分别测试了并发数1、8、16、32、64五个档位。EFS和FSx本身是多通道协议，所以主要看挂载参数和文件系统吞吐设置的组合。并发数不是越大越好，超过一定阈值后，反而会因为请求排队、CPU上下文切换而下降。具体怎么找这个阈值，我会在3.3里讲一个推导方法。

数据采集我做得比较细。每轮测试记录开始时间、结束时间、文件总数、总大小、成功读取文件数、失败重试次数、峰值内存、峰值CPU。每轮重复至少三次，取中位数作为结果，避免单次网络波动或者限流导致的数据失真。最终所有结果连同环境信息、脚本版本、日期一起写入CSV，方便复盘。

3. 基准测试实操：脚本、命令与参数推导

3.1 三步走脚本：准备、计时与数据落盘

基准测试脚本我写成了三个阶段的独立脚本，每个阶段可以单独运行，互不依赖，这样方便排查。

第一阶段是准备脚本，负责生成Checkpoint样本。为了模拟分布式场景，我起多个进程，每个进程保存自己的shard文件。这里有个小细节：文件内容必须用真实的随机张量，不能用全零数据，否则后续校验阶段会“失真”。准备脚本不参与任何计时，生成完就退出。

第二阶段是读取计时脚本，核心逻辑非常简单：

import time import glob import hashlib import os def compute_md5(file_path): h = hashlib.md5() with open(file_path, "rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): h.update(chunk) return h.hexdigest() file_list = glob.glob("/mnt/checkpoint/*") start = time.time() total_size = 0 success = 0 for fp in file_list: # 模拟框架读取：这里直接用hash校验代替反序列化，避免干扰 compute_md5(fp) total_size += os.path.getsize(fp) success += 1 elapsed = time.time() - start throughput = total_size / elapsed / 1024 / 1024 / 1024 print(f"elapsed={elapsed:.2f}s throughput={throughput:.2f}GB/s success={success}/{len(file_list)}")

用MD5校验来模拟读取，是因为它不会在内存里重建巨大的张量结构，能够更纯粹地测存储和网络能力。但注意，这只代表“读取”这个环节，真实的恢复时间还要加第三阶段。

第三阶段是用PyTorch真实加载Checkpoint并计时。对于分布式Checkpoint，我会调用torch.distributed.checkpoint.load，并确保加载后的参数量和预期一致。加载完成后，再跑一个极简的前向步骤，确认模型能够正常执行，这个步骤用来判断恢复是否真成功，而不是仅仅文件读到了本地。

脚本里我还会在每轮测试前记录一次环境变量：

echo "=== ENV ==="; nvidia-smi; df -h /mnt; ulimit -n; sysctl net.core.somaxconn

这些信息对后期排查性能问题特别有用。ulimit -n如果不调大，文件一多很容易报“Too many open files”。

3.2 S3、EFS、EBS与FSx挂载要点

存储方案不同，命令和参数完全不一样，每一个都有值得记下来的坑。

S3这边，我优先用s5cmd而不是AWS CLI。AWS CLI的aws s3 cp是单线程的，拉30GB文件能让你等到怀疑人生。s5cmd支持并发下载，命令也很简单：

# 并发32下载整个prefix s5cmd --numworkers 32 cp "s3://bucket/checkpoint/run-001/" /mnt/checkpoint/

也试过mountpoint-s3，它能把S3桶挂载成本地目录，读取时自动并行。不过要注意，mountpoint-s3默认是只读挂载，适合直接给框架读文件，但不适合作为Checkpoint的写入目录。它最大的优势是对应用透明，你不需要在代码里区分“这是S3还是本地路径”。

EFS的挂载参数会影响吞吐。我用的最小参数组合是这样：

sudo mount -t nfs4 -o rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport -o tls fs-xxxx.efs.us-east-1.amazonaws.com:/ /mnt/efs

rsize和wsize默认值往往偏小，手动改成1MB能显著提升大文件顺序读性能。noresvport这个参数在NFS重连时有用，减少连接复用导致的“卡住”问题。

EBS本质是块设备，读取快慢取决于卷类型和IOPS设置。gp3卷可以独立调IOPS和吞吐，测试时我把IOPS调到10000，吞吐调到500MB/s，再在卷上创建XFS文件系统。如果只是测试热读，直接把Checkpoint放在已经挂载好的EBS上即可，无需额外操作。

FSx for Lustre相对省心，创建时关联同一个S3桶，挂载后就可以按路径读写。要注意的是默认情况下FSx的导入导出是Lazy方式，也就是说你先要从S3load数据到文件系统，再读才有速度。实际操作里，我会先用lustre客户端做一次预热导入，再开始计时。命令大致是：

sudo lfs hsm_restore -r /mnt/fsx/checkpoint/

这一步不做好，第一次读取可能走的是S3后端的懒加载，测出来的是S3的性能，不是FSx的性能，数据会失真。

3.3 理论带宽与并发数推导：先算上限再跑测试

在跑任何基准之前，建议先做一个理论估算，心里有个“不可能超过多少”的底。以medium规格（30GB）为例，假设实例网络能跑到10GB/s（实际允许值），那从任何存储拉完数据的物理极限就是3秒。但实际肯定做不到，因为存储服务不可能让单个实例占满全部带宽，还会受到每连接带宽、请求数、文件系统元数据等限制。

我的做法是先做小样本探测。随便从一个S3 prefix拉一个1GB的大文件，记录单线程、8线程、16线程、32线程各自的吞吐。比如实测结果：

可以看到，16线程以后吞吐增长趋缓，32线程已经接近瓶颈。这时候把并发定在16到24之间，既不会浪费CPU，又留了余量应对波动。如果你的文件数量很多，并发不能只看带宽，还要算QPS。S3单前缀的GET请求速率大约每秒几千次，5000个小文件全量下载，光请求就要好几秒。这种情况下，优先考虑压缩文件数量（打包成tar），比单纯加大并发有效得多。

FSx和EFS的推导逻辑类似，但它们的上限通常由文件系统吞吐和实例网卡共同决定，所以并发的影响没有S3那么明显。EBS的话，因为底层就是块存储，顺序读性能相对稳定，主要限制是卷的吞吐上限。

4. 实测数据解读与成本分析

4.1 同环境下不同存储方案的结果对比

下面这组数据来自我当前测试环境的实测，只代表本次环境下的表现。你用自己的账号在另一个时间跑，数值可能完全不同，但对比关系和趋势有参考价值。

几个关键发现：S3在medium规格下能跑出1.4GB/s的吞吐，说明方向对了；但到了large规格，吞吐掉到0.85GB/s，主要原因是5000个文件导致请求排队。FSx for Lustre在两种规格下表现都稳定，基本维持在3.5GB/s左右，说明它对小文件混合场景更友好。EBS热读最快，但它是本地盘，机器没了数据也没了，只适合进程崩溃恢复的场景。

再说一个意外：EFS的冷读成绩比我想象中差，尽管设置了弹性吞吐，但1000个小文件的元数据操作拖累了整体表现。如果你用EFS保存大量小文件，性能大概率不理想，除非改用更少、更大的文件。

4.2 恢复时间拆分：瓶颈到底在哪一段

只看总耗时是不够的，我每次都把恢复时间拆成三段：读取阶段、校验阶段、加载阶段。以medium规格S3冷读那组来说，实测如下：

• 读取阶段：21.4秒（占55%）
• 校验阶段：7.1秒（占18%）
• 加载阶段：10.2秒（占27%）
• 总恢复时间：38.7秒

读取阶段占大头，这大家都知道。但很多人没注意的是校验阶段——如果你用hash校验每个文件，文件多且小的时候，这个阶段会很痛。我甚至见过有人在上万个文件的场景里，校验阶段比读取阶段还长。所以，如果你的RTO极其敏感，可以考虑去掉全量校验，改成抽样校验，或者用safetensors格式的header校验，只校验元数据，速度会快很多。

加载阶段也值得优化。torch.load默认要先把整个对象load到CPU内存，再拷贝到GPU，这个过程有大量的内存分配和反序列化开销。对于分布式Checkpoint，加载时还涉及reshard，把各rank的shard重新分配到当前可用GPU上，这部分逻辑如果写的不好，耗时会翻倍。优化手段包括：用mmap方式加载模型权重、使用safetensors避免pickle开销、加载前用meta device初始化模型再赋值参数，减少峰值CPU内存和拷贝。

4.3 快速恢复能省多少钱：一笔账算清RTO价值

RTO不只是技术指标，它直接换算成钱。假设一个8卡A100实例，按量价格约每小时32美元，Spot价按60%折扣算，大约13美元一小时。但注意，Spot价格波动大，这里是示意值。

假设一次训练中断需要重新拉起实例并恢复Checkpoint。如果恢复需要40分钟，Spot实例这40分钟的租金是8.7美元；如果通过优化存储和加载方式，恢复时间缩短到15分钟，租金变成3.25美元。一次恢复省下5.45美元。如果一天平均被回收3次，一天省16.35美元，一个月省差不多490美元。

这还不算最关键的隐形成本——机器空转。恢复期间GPU完全空闲，集群越大损失越大。一个8卡实例每小时32美元，空转40分钟意味着21.3美元的算力浪费。所以，把恢复时间从40分钟压缩到15分钟，不只是省租金，更重要的是把宝贵的GPU时间用在训练上。搞训练的人都知道，GPU空转一秒钟都肉疼。

5. Checkpoint读取恢复常见问题与排查实录

5.1 读取慢的五个排查点，照着查就行

如果你遇到Checkpoint读取慢，先别急着骂存储，按下面顺序排查一遍。

第一，并发数够不够。S3用aws s3 cp默认单线程，慢是正常的，换成s5cmd或者调大--numworkers。EFS则检查挂载参数和文件系统吞吐设置。

第二，文件数量是否过多。当单个文件平均小于几百KB，请求数就会成为瓶颈。最有效的解法是把一堆小文件打包成一个大tar包或者合并成更少的分片文件。

第三，网络路径是否绕路。确认实例是否通过VPC Gateway Endpoint访问S3，而不是走公网。跨可用区访问EFS和FSx也会显著增加延迟。检查方法很简单，看挂载或Endpoint配置。

第四，挂载参数是否正确。EFS的rsize和wsize过小会严重限制吞吐，建议设置成1MB以上。

第五，检查安全软件和系统限制。比如默认的ulimit -n太小，文件一多直接报错；或者实例上的监控代理在扫描文件目录，干扰了读取。

5.2 恢复时的OOM与设备映射问题

Checkpoint恢复最常见的一个坑是OOM。模型权重本身可能只有14GB，但加上Adam优化器状态，整个Checkpoint膨胀到50GB以上，如果再通过torch.load一次性加载到CPU内存，普通实例的CPU内存根本扛不住。解决办法是用torch.load(..., mmap=True)或者safetensors的load_file，它们支持按需映射文件到内存，不一次性把整个对象塞进来。

第二个坑是设备映射。分布式训练保存Checkpoint时，每个rank保存的是自己的shard，但恢复时可能GPU的数量、拓扑都没变，也可能变了。比如你原来用8卡训练，恢复时只有4卡可用，这时候就要做reshard。torch.distributed.checkpoint的load会自动处理不同world size之间的shard重新分配，但用错了API，比如用普通的load_state_dict去加载分布式保存的Checkpoint，就会报shape不匹配或者参数找不到。我的建议是：训练保存和恢复尽量用同一套分布式Checkpoint API，不要混用。

5.3 小文件地狱：为什么需要打包策略

5000个文件在large规格里已经让人头疼，如果Checkpoint是上万个几KB的metadata文件，那纯粹是灾难。读取这类文件时，大部分时间都耗在文件打开、目录遍历、S3请求上，而不是数据传输。

打包策略并不复杂。把一堆小文件合并成少量大文件，比如打成tar包，然后用流式读取。PyTorch生态里已经有一些做法，比如将模型权重用单个二进制文件保存，并将元数据放在文件头部；或者用tar把所有文件打成一个压缩包，下载后再解包。不过要权衡解包带来的CPU开销和时间。在medium规格里，我把1000个小文件打包成4个大文件，S3读取时间从21.4秒降到14.6秒，校验时间也大幅缩短。代价是打包和解包各花3秒左右，整体还是划算的。

如果你的训练运行时间长、恢复频率高，建议直接把Checkpoint的存储格式设计成“少量大文件”的结构，而不是依赖事后打包。这个设计越早做，后期越省心。

5.4 把基准测试做成训练基础设施的一部分

很多团队做一次基准测试就完事了，我觉得不够。环境会变，框架会升级，数据规模会变，正确做法是把这套基准测试脚本纳入CI流程，或者至少在每个训练项目启动前跑一遍。具体来说，我有三条习惯，分享给大家。

一是任何涉及框架版本、实例类型、存储类型变动的操作，都要重跑一遍基准。不要相信“升级应该不影响性能”这种话，有一次我把PyTorch从2.0升到2.1，torch.save默认格式没变，但分布式Checkpoint的加载实现变了，恢复时间莫名多了20%，只有跑基准才能暴露。

二是训练中断恢复演练要定期做。不是等到故障发生了才想起恢复，而是主动模拟一次Spot回收，计时、记录日志、复盘问题。团队里如果有新人接手，这也是一份很好的SOP。

三是把基准测试结果和成本账单挂在一起看。你优化了多少秒RTO，省了多少钱，这些数字是给管理层和财务看的。技术指标要翻译成业务语言，才能让更多同事支持你把时间投入到这类“不直接产生模型效果”的基础设施建设上。

我个人在实际操作中的体会是，Checkpoint恢复基准不是一次性的项目，而是一种习惯。每次看到训练任务安稳地跑过之前容易中断的节点，我都会想起那些躲在脚本里的计时器和曲线。下次如果你也被某个读取慢、恢复失败的问题折磨到凌晨，欢迎回来翻这篇，照着5.1的排查清单过一遍，大概率就能定位到问题所在。或者你发现自己遇到的情况比我写的更复杂，也可以沿着这条思路自己设计一套更贴近业务场景的基准，数据永远比感觉靠谱。