如果你的训练任务在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还要做跨进程的reshardDeepSpeed的Zero-3状态更大动辄几十GB。就算是同一个模型用不同版本框架保存出来的Checkpoint结构也有差异恢复性能自然不同。第三网络路径的影响。数据从S3走公网Endpoint还是VPC Gateway Endpoint跨可用区还是同可用区延迟和带宽的表现都不一样。我们曾经因为没走Gateway Endpoint从S3拉数据额外多花了两倍时间排查了半天才意识到是网络路径的问题。所以别省这一步。把基准测试当成训练基础设施的一部分每次改环境、升版本、换存储之后都跑一遍用数据替代直觉。2. 测试环境与基准设计先搭台子再跑分2.1 实例、存储与框架版本怎么选测试环境的搭建要尽量贴近真实训练环境否则测出来的数据没有参考价值。我这次使用的实例是p4d.24xlarge8张A100 40GB网络带宽400Gbps约合51.2GB/s的理论上限。不过要注意这个400Gbps是实例级别的总带宽实际到单个S3桶、单个EFS文件系统时不可能吃满能跑到几GB/s已经算不错。存储方面我准备了四条路线S3标准存储桶使用VPC Gateway Endpoint、EFS选择弹性吞吐模式置备了足够的吞吐量、EBSgp33000 IOPS基线125MB/s基线吞吐测试时开到10000 IOPS和500MB/s吞吐、FSx for Lustre做成了与S3关联的数据仓库模式。每个存储方案都单独测试避免互相干扰。软件版本锁定非常关键。PyTorch选2.1.2DeepSpeed选0.14.4Hugging Face Transformers选4.39.3Python版本3.10。为什么版本要锁得这么死因为DeepSpeed的ZeRO-3 Checkpoint格式在0.13到0.14之间有过不兼容的改动Hugging Face的模型结构序列化方式也可能变化。你不想辛辛苦苦测出的结果过两个月因为框架升级就全部失效。我的做法是每次测试都记录一份pip freeze到实验日志里。2.2 模拟Checkpoint文件结构与三种测试规格真实的大模型Checkpoint不是单个文件。以7B参数模型为例混合精度训练下模型权重大约14GBBF16Adam优化器状态至少还要两倍也就是28GB往上整体一个Checkpoint接近50GB甚至更多。而且分布式训练时每个GPU rank会各自保存自己分管的那部分权重和优化器状态再加上tokenizer、config、调度器状态等小文件整体文件数量很容易上千。为了做基准测试我构造了三种规格的Checkpoint样本都是模拟真实结构生成的规格文件总大小文件数量典型来源small5GB100单卡微调关闭优化器状态medium30GB10007B模型混合精度微调含Adam状态large100GB5000更大模型或更多分片文件更碎生成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(felapsed{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 rsize1048576,wsize1048576,hard,timeo600,retrans2,noresvport -o tls fs-xxxx.efs.us-east-1.amazonaws.com:/ /mnt/efsrsize和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线程各自的吞吐。比如实测结果并发数拉取1GB单文件耗时折合吞吐18.2s0.12GB/s81.9s0.53GB/s161.2s0.83GB/s321.1s0.91GB/s可以看到16线程以后吞吐增长趋缓32线程已经接近瓶颈。这时候把并发定在16到24之间既不会浪费CPU又留了余量应对波动。如果你的文件数量很多并发不能只看带宽还要算QPS。S3单前缀的GET请求速率大约每秒几千次5000个小文件全量下载光请求就要好几秒。这种情况下优先考虑压缩文件数量打包成tar比单纯加大并发有效得多。FSx和EFS的推导逻辑类似但它们的上限通常由文件系统吞吐和实例网卡共同决定所以并发的影响没有S3那么明显。EBS的话因为底层就是块存储顺序读性能相对稳定主要限制是卷的吞吐上限。4. 实测数据解读与成本分析4.1 同环境下不同存储方案的结果对比下面这组数据来自我当前测试环境的实测只代表本次环境下的表现。你用自己的账号在另一个时间跑数值可能完全不同但对比关系和趋势有参考价值。存储方案规格场景耗时吞吐实际恢复时间含加载S3s5cmd并发32medium 30GB/1000文件冷读21.4s1.40GB/s38.7sS3mountpoint-s3medium 30GB/1000文件冷读23.1s1.30GB/s40.2sEFS弹性吞吐medium 30GB/1000文件冷读57.6s0.52GB/s76.5sEBSgp3 500MB/smedium 30GB/1000文件热读3.8s7.89GB/s12.4sFSx for Lustremedium 30GB/1000文件冷读8.2s3.66GB/s19.5sS3large 100GB/5000文件冷读118.3s0.85GB/s151.2sFSx for Lustrelarge 100GB/5000文件冷读28.6s3.50GB/s55.4s几个关键发现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(..., mmapTrue)或者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.1torch.save默认格式没变但分布式Checkpoint的加载实现变了恢复时间莫名多了20%只有跑基准才能暴露。二是训练中断恢复演练要定期做。不是等到故障发生了才想起恢复而是主动模拟一次Spot回收计时、记录日志、复盘问题。团队里如果有新人接手这也是一份很好的SOP。三是把基准测试结果和成本账单挂在一起看。你优化了多少秒RTO省了多少钱这些数字是给管理层和财务看的。技术指标要翻译成业务语言才能让更多同事支持你把时间投入到这类“不直接产生模型效果”的基础设施建设上。我个人在实际操作中的体会是Checkpoint恢复基准不是一次性的项目而是一种习惯。每次看到训练任务安稳地跑过之前容易中断的节点我都会想起那些躲在脚本里的计时器和曲线。下次如果你也被某个读取慢、恢复失败的问题折磨到凌晨欢迎回来翻这篇照着5.1的排查清单过一遍大概率就能定位到问题所在。或者你发现自己遇到的情况比我写的更复杂也可以沿着这条思路自己设计一套更贴近业务场景的基准数据永远比感觉靠谱。