示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载FastPersist 是 DeepSpeed 提出的、利用 NVMe 存储加速模型检查点checkpoint保存的优化技术论文出处arXiv:2406.13768。本指南以 deepnvme/model_checkpoint/README.md 为脉络结合该目录下完整的微基准测试源码讲解 FastPersist 的启用方式、三个微基准的用法与全部命令行参数、输出日志的解读方法并深入到 torch 序列化补丁与底层 I/O 写入器的实现原理帮助读者在自己的环境中复现从 0.69 GB/s 到 17.75 GB/s量级的检查点吞吐提升。一、FastPersist 是什么用 NVMe 接管 torch 序列化大模型训练中检查点保存往往成为训练吞吐的隐形瓶颈模型权重与优化器状态动辄十几 GB用传统方式逐块写入磁盘时GPU 需要长时间等待 I/O造成计算空窗。FastPersist 的核心思路是把 torch 序列化产生的数据流直接导向高性能 NVMe 存储利用 DeepSpeed 的异步 I/OAsyncIO即ds_aio与 NVIDIA GPUDirect StorageGDS能力让检查点写入不再阻塞训练主流程。从目录结构看deepnvme/model_checkpoint 集中提供了三类验证素材素材文件作用torch 序列化补丁torch/serialization_orig_v2.6.0.py、torch/serialization_fast_v2.6.0.py原始版与补丁版serialization.py用于替换 torch 安装目录中的同名文件微基准脚本torch_save_tensor.py、torch_save_model.py、deepspeed_save_model.py分别测试裸张量、HF 模型torch.save 集成、HF 模型DeepSpeed 集成的保存性能公共工具torch_save_utils.py、save_model_utils.py五种写入器封装、AIO/GDS handle 构造、模型与参数解析FastPersist 设计上与 torch 检查点机制集成官方在 torch 2.6.0 版本上完成验证。集成方式很轻只需要对 torch 的序列化文件做一处小修改因此仓库同时提供了原始版与补丁版两个serialization.py供对照。二、启用 FastPersist替换 torch 的 serialization.pyFastPersist 的启用步骤在 README 中有明确说明将补丁版 torch/serialization_fast_v2.6.0.py 覆盖到当前 torch 安装目录下的torch/serialization.py即# 将仓库中的补丁版序列化文件覆盖到 torch 安装路径 cp deepnvme/model_checkpoint/torch/serialization_fast_v2.6.0.py \ $(python -c import torch, os; print(os.path.dirname(torch.__file__)))/serialization.py覆盖前建议备份原始文件仓库中的 serialization_orig_v2.6.0.py 即是原版备份。该操作要求 torch 版本为2.6.0其他版本可能因源码差异导致补丁不完全匹配这是 README 明确给出的适用前提。补丁到底改了什么通过对比两个文件可以精确定位改动集中在_legacy_save函数的 storage 写出环节。原始版本是这样逐个写 storage 的for key in serialized_storage_keys: storage, dtype serialized_storages[key] storage._write_file( f, _should_read_directly(f), True, torch._utils._element_size(dtype) )补丁版则先探测写入对象是否实现了save_torch_storage_object_list接口if hasattr(f, save_torch_storage_object_list): sorted_storage_objects [serialized_storages[key] for key in serialized_storage_keys] f.save_torch_storage_object_list(sorted_storage_objects, True) else: # 回退到原始的逐 storage 写出逻辑 ...这一小段改动是整个 FastPersist 机制的接入点当torch.save(fwriter, ...)的f是 DeepSpeed 提供的FastFileWriter时所有序列化 storage 会一次性交给写入器的save_torch_storage_object_list批量下盘当f仍是普通文件对象时走原始逐块写出路径行为完全不变。也就是说补丁对 vanillatorch.save()无任何副作用只增加了对 FastPersist 写入器的识别与分发。三、三个微基准脚本定位与统一输出格式目录提供三个单进程微基准都通过python直接启动均支持--help查看全部配置项torch_save_tensor.py将裸 PyTorch 张量经torch.save()写盘用于在最小粒度上隔离序列化与 I/O 开销torch_save_model.py将 HuggingFace 模型经torch.save()集成写盘贴近实际权重落盘场景deepspeed_save_model.py将 HF 模型交给 DeepSpeedsave_checkpoint()集成写盘覆盖 ZeRO 优化器状态等完整训练态。三个脚本统一维护一组标签 → 写入器的映射见 torch_save_utils.py依次测出五种机制标签底层写入器说明test_savevanillatorch.save()基准对照常规文件直写test_ds_mock_saveMockFileWriter模拟写入便于估算理论开销test_ds_py_savePyFileWriterDeepSpeed 纯 Python 写入器test_ds_aio_fast_saveFastFileWriter AsyncIOFastPersist CPU 钉页缓冲bounce buffertest_ds_gds_fast_saveFastFileWriter GDSFastPersist NVIDIA GPUDirect Storage其中test_ds_gds_fast_save依赖 GPU 上的钉页张量torch_save_utils.py因此在三个脚本中当未指定--gpu时该测试会被自动跳过if tag test_ds_gds_fast_save and not args.gpu: continue所有脚本的最终输出格式统一为tag -- 大小 GB, 耗时 secs, 吞吐 GB/sREADME 中给出的 8 块 PCIe Gen4 NVMe RAID-0 实测示例即为该格式的标准产物test_save -- 14.23 GB, 20.72 secs, 0.69 GB/s test_ds_aio_fast_save -- 14.23 GB, 0.80 secs, 17.75 GB/s同一份 14.23 GB 检查点vanillatorch.save()耗时 20.72 秒0.69 GB/sFastPersist CPU bounce buffer 仅耗时 0.80 秒17.75 GB/s吞吐提升约 25 倍。注意这是仓库作者在特定硬件RAID-0 条带化 8 盘 NVMe上收集的示例数据实际收益会随存储介质、队列深度、缓冲大小等因素变化。四、微基准一裸张量保存torch_save_tensor.pytorch_save_tensor.py 用于测量纯张量的保存性能。核心流程是按--mb_size生成一个torch.uint8张量指定--gpu时驻留 GPU随后依次用五种写入器计时写盘。其命令行参数torch_save_tensor.py参数默认值说明--folder必填I/O 目标目录建议指向 NVMe 挂载点--mb_size必填张量大小MB--zipfile关使用 torch zipfile 保存格式--gpu关使用 GPU 张量启用后才会跑 GDS 测试--io_buffer_mb64钉页 I/O 缓冲大小MB--no-statistics关抑制底层性能统计输出--single_io_buffer关禁用 I/O 缓冲双缓冲运行示例把 2 GB 张量写入 NVMe 目录python torch_save_tensor.py --folder /mnt/nvme0 --mb_size 2048 --gpu值得注意的是张量生成方式数据通过torch.randint(high128, ...)构造与真实模型权重一样非零避免让压缩类 I/O 路径取巧。输出文件按{tag}_{mb_size}MB.pt命名例如test_save_2048MB.pt、test_ds_aio_fast_save_2048MB.pt。五、微基准二HF 模型经 torch.save() 集成torch_save_model.pytorch_save_model.py 把 HF 模型封装进{model: model}或{model: model, optimizer: optimizer}的字典后再交给torch.save()以此验证 FastPersist 对完整模型态含优化器状态的加速。支持的模型标签定义在 save_model_utils.py 的HF_MODELS_DICT中标签HF 模型名tiny-t5hf-internal-testing/tiny-random-t5T5 生成模型gpt2-largegpt2-largephi3microsoft/Phi-3.5-mini-instructphi3-vmicrosoft/Phi-3.5-vision-instructllama3-1Bmeta-llama/Llama-3.2-1B除tiny-t5走T5ForConditionalGeneration外其余模型均通过AutoModelForCausalLM加载同时加载对应分词器。README 中的示范命令python torch_save_model.py --model phi3 --folder /mnt/nvme0 --gpu完整参数save_model_utils.py参数默认值说明--folder必填I/O 目录--model必填上表任一模型标签--local_rank0本地 rank--zipfile关使用 torch zipfile 保存格式--optimizer关检查点中包含优化器状态--fused关使用 fused FP16 优化器--gpu关使用 GPU 张量--half关使用半精度张量--io_buffer_mb64钉页 I/O 缓冲MB--zero_stage0ZeRO 优化阶段本脚本内部未实际使用保留给 DeepSpeed 集成--cpu_offload关启用优化器状态 CPU offload--no-statistics关抑制底层性能统计--single_io_buffer关禁用双缓冲--safetensors关使用 safetensors 读写--regular_torch_save关使用 vanilla torch.save注意两个参数联动约束见 save_model_utils.py 的validate_arguments--model必须是上表标签--optimizer与--half同时开启时必须配合--gpu因为混合精度优化器状态只支持 GPU 张量。优化器默认走FP16_Wrapper封装--fused时选 fused 版本并在首次step()前把所有参数梯度置零来初始化优化器状态torch_save_model.py保证测出的检查点体积包含真实的优化器状态数据。六、微基准三DeepSpeed 集成deepspeed_save_model.pydeepspeed_save_model.py 走完整 DeepSpeed 引擎路径通过deepspeed.initialize()构造引擎deepspeed_save_model.py再调用ds_engine.save_checkpoint(save_dirfolder, tagtag)落盘检查点按{folder}/{ckpt_name}/{tag}/目录组织测完通过_get_folder_size累加目录内所有文件得到总体积。与脚本二不同这里的关键参数通过 DeepSpeed 配置文件注入deepspeed_save_model.py{ zero_optimization: { stage: 0, cpu_offload: false }, fp16: { enabled: false }, optimizer: { type: Adam, params: { torch_adam: true } }, checkpoint: { checkpoint_serialization: false, writer: { type: fast, io_buffer_size: 67108864, io_buffer_double: true, show_statistics: true, data_parallel: socket } }, aio: { block_size: 8388608, queue_depth: 8, single_submit: false, overlap_events: true, intra_op_parallelism: 2, use_gds: false } }配置项随--zero_stage、--cpu_offload、--half、--zipfile、--io_buffer_mb、--single_io_buffer等命令行参数动态生成writer.type则按标签映射为mock/python/fastdeepspeed_save_model.pyuse_gds仅在fast GDS 标签时为真。该脚本还内建了单进程分布式初始化deepspeed_save_model.py通过设置MASTER_ADDRlocalhost、WORLD_SIZE1等环境变量并以 rank 0 初始化 NCCL 进程组让单机单卡也能跑通 DeepSpeed 引擎。运行方式与其他脚本一致python deepspeed_save_model.py --model phi3 --folder /mnt/nvme0 --gpu --optimizer七、五种写入器的底层实现从 bounce buffer 到 GDS所有写入器的差异都封装在 torch_save_utils.py 中理解它们即可理解 FastPersist 的性能来源。AIO/GDS handle 构造。FastPersist 路径基于 DeepSpeed 的AsyncIOBuilder与GDSBuilder动态构建底层句柄torch_save_utils.pyh AsyncIOBuilder().load().aio_handle( block_size8 * (1024**2), # 8 MB 块 queue_depth8, # 异步队列深度 single_submitFalse, # 关闭单提交模式 overlap_eventsFalse, # 事件重叠开关 intra_op_parallelism1) # 操作内并行度脚本启动时load_io_ops会先做兼容性探测torch_save_utils.pyAsyncIO 与 GDS 扩展只有在构建器is_compatible()通过时才加载。钉页缓冲bounce buffer。CPU 路径的缓冲是一个钉页pinned memory张量默认 64 MBPINNED_BUFFER_MB可经--io_buffer_mb调整GDS 路径则是在 GPU 设备上分配张量并调用h.pin_device_tensor()钉住torch_save_utils.py。CPU bounce buffer 的作用是充当 GPU 显存与 NVMe 之间的 DMA 中转站——数据先拷入钉页缓冲再由 AIO 异步引擎零拷贝写入存储GDS 则更进一步允许 GPU 显存与 NVMe 之间直通绕过 CPU 内存中转。FastFileWriter 装配。两种 FastPersist 路径共用同一个核心函数torch_save_utils.py差异只在句柄与缓冲的来源fast_writer_config FastFileWriterConfig( dnvme_handleh, # AIO 或 GDS 句柄 pinned_tensorpinned_memory, # CPU 钉页或 GPU 钉页缓冲 double_buffernot args.single_io_buffer, # 默认双缓冲 num_parallel_writers1, writer_rank0) ds_fast_writer FastFileWriter(file_pathfile, configfast_writer_config) torch.save(fds_fast_writer, objbuffer, ...) ds_fast_writer.close() # 强制 flush 到存储close()之后的_dump_state()默认开启--no-statistics可关闭会打印底层 I/O 统计。双缓冲io_buffer_doubleTrue/ 不传--single_io_buffer让 DMA 填充与异步写出两个阶段重叠是吞吐最大化的重要开关。vanilla 对照。test_save直接调用torch.save(ffile, ...)MockFileWriter与PyFileWriter则是 DeepSpeed 提供的另外两种torch.save()集成写入器用于在未启用补丁/未加载 AIO时给出基准线。八、加载验证torch_save_load_model.py 的闭环测试目录中还有一个不在 README 编号之列但值得关注的辅助脚本 torch_save_load_model.py它在保存之后立即用torch.load(..., weights_onlyTrue)或 safetensors把权重读回并用torch.allclose(pre_logits, post_logits, atol1e-3, rtol1e-3)断言加载前后模型前向输出一致torch_save_load_model.py。该脚本支持--regular_torch_savevanilla 对照与--safetensorssafetensors 格式两种旁路默认走_test_ds_fast_save的 FastPersist 路径。它可以作为FastPersist 落盘数据与原生 torch 保存等价的正确性闭环验证手段。九、复现建议与运行前提综合 README 与源码复现 FastPersist 收益的完整路径为确认硬件至少一块 NVMe SSD推荐多盘 RAID-0 以获得 README 示例量级的带宽GDS 路径需要 NVIDIA GPU 与支持 GDS 的驱动/存储环境确认软件torch 2.6.0、DeepSpeed提供AsyncIOBuilder、GDSBuilder、FastFileWriter等组件、transformersrequirements.txt 声明脚本依赖的 DeepSpeed 扩展需能正常加载打补丁用 serialization_fast_v2.6.0.py 覆盖 torch 安装目录下的torch/serialization.py先备份原版跑基准按上文命令依次执行三个脚本--folder指向 NVMe 挂载点测模型时--gpu以获得含 GDS 的完整结果读日志在输出中搜索test_save、test_ds_aio_fast_save、test_ds_gds_fast_save等行对比各机制吞吐test_ds_*系列默认还会附带_dump_state()的底层统计可进一步观察队列深度、单次 I/O 大小等细节。最后提醒两点边界README 中的性能数字是在特定硬件8 盘 PCIe Gen4 NVMe RAID-0上采集的示例不代表任意环境下的保证值补丁文件针对 torch 2.6.0 编写升级 torch 版本前需要重新评估补丁兼容性。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐DeepSpeed 模型检查点Model Checkpointing实战指南保存、加载与 ZeRO fp32 权重恢复DeepSpeed 模型检查点Model Checkpointing实战指南保存、加载与 ZeRO fp32 权重恢复 本指南以 DeepSpeed 官方推理引擎大模型Transformers 与 DeepSpeed 集成实战ZeRO 分阶段训练、CPU/NVMe 卸载与大规模模型推理指南Transformers 与 DeepSpeed 集成实战ZeRO 分阶段训练、CPU/NVMe 卸载与大规模模型推理指南 导读 本文基于 Hugging F人工智能大模型深度学习NLP预训练微调模型推理服务mambaout_base_wide_rw.sw_e500_in1k完全指南从零开始的图像分类实战教程mambaout_base_wide_rw.sw_e500_in1k完全指南从零开始的图像分类实战教程 想要快速上手高效的图像分类模型吗mambaout_b上一篇5分钟精通暗黑2存档编辑Diablo Edit2终极角色定制解决方案下一篇Operit 市场溯源机制ToolPkg 发布注入、导入校验与来源展示的完整实现创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考