企业官网建设流程全解析

• 示例工程

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/O（AsyncIO，即ds_aio）与 NVIDIA GPUDirect Storage（GDS）能力，让检查点写入不再阻塞训练主流程。

从目录结构看，deepnvme/model_checkpoint 集中提供了三类验证素材：

FastPersist 设计上与 torch 检查点机制集成，官方在 torch 2.6.0 版本上完成验证。集成方式很"轻"：只需要对 torch 的序列化文件做一处小修改，因此仓库同时提供了原始版与补丁版两个serialization.py供对照。

二、启用 FastPersist：替换 torch 的 serialization.py

FastPersist 的启用步骤在 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(f=writer, ...)的f是 DeepSpeed 提供的FastFileWriter时，所有序列化 storage 会一次性交给写入器的save_torch_storage_object_list批量下盘；当f仍是普通文件对象时，走原始逐块写出路径，行为完全不变。也就是说，补丁对 vanillatorch.save()无任何副作用，只增加了对 FastPersist 写入器的识别与分发。

三、三个微基准：脚本定位与统一输出格式

目录提供三个单进程微基准，都通过python直接启动，均支持--help查看全部配置项：

1. torch_save_tensor.py：将裸 PyTorch 张量经torch.save()写盘，用于在最小粒度上隔离序列化与 I/O 开销；
2. torch_save_model.py：将 HuggingFace 模型经torch.save()集成写盘，贴近实际"权重落盘"场景；
3. deepspeed_save_model.py：将 HF 模型交给 DeepSpeedsave_checkpoint()集成写盘，覆盖 ZeRO 优化器状态等完整训练态。

三个脚本统一维护一组"标签 → 写入器"的映射（见 torch_save_utils.py），依次测出五种机制：

其中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/s

README 中给出的 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/s），FastPersist + CPU bounce buffer 仅耗时 0.80 秒（17.75 GB/s），吞吐提升约 25 倍。注意：这是仓库作者在特定硬件（RAID-0 条带化 8 盘 NVMe）上收集的示例数据，实际收益会随存储介质、队列深度、缓冲大小等因素变化。

四、微基准一：裸张量保存（torch_save_tensor.py）

torch_save_tensor.py 用于测量纯张量的保存性能。核心流程是：按--mb_size生成一个torch.uint8张量（指定--gpu时驻留 GPU），随后依次用五种写入器计时写盘。

其命令行参数（torch_save_tensor.py）：

运行示例（把 2 GB 张量写入 NVMe 目录）：

python torch_save_tensor.py --folder /mnt/nvme0 --mb_size 2048 --gpu

值得注意的是张量生成方式：数据通过torch.randint(high=128, ...)构造，与真实模型权重一样非零，避免让压缩类 I/O 路径"取巧"。输出文件按{tag}_{mb_size}MB.pt命名，例如test_save_2048MB.pt、test_ds_aio_fast_save_2048MB.pt。

五、微基准二：HF 模型经 torch.save() 集成（torch_save_model.py）

torch_save_model.py 把 HF 模型封装进{'model': model}或{'model': model, 'optimizer': optimizer}的字典后，再交给torch.save()，以此验证 FastPersist 对完整模型态（含优化器状态）的加速。

支持的模型标签定义在 save_model_utils.py 的HF_MODELS_DICT中：

除tiny-t5走T5ForConditionalGeneration外，其余模型均通过AutoModelForCausalLM加载，同时加载对应分词器。README 中的示范命令：

python torch_save_model.py --model phi3 --folder /mnt/nvme0 --gpu

完整参数（save_model_utils.py）：

注意两个参数联动约束（见 save_model_utils.py 的validate_arguments）：--model必须是上表标签；--optimizer与--half同时开启时必须配合--gpu，因为混合精度优化器状态只支持 GPU 张量。优化器默认走FP16_Wrapper封装（--fused时选 fused 版本），并在首次step()前把所有参数梯度置零来初始化优化器状态（torch_save_model.py），保证测出的检查点体积包含真实的优化器状态数据。

六、微基准三：DeepSpeed 集成（deepspeed_save_model.py）

deepspeed_save_model.py 走完整 DeepSpeed 引擎路径：通过deepspeed.initialize()构造引擎（deepspeed_save_model.py），再调用ds_engine.save_checkpoint(save_dir=folder, tag=tag)落盘，检查点按{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/fast（deepspeed_save_model.py），use_gds仅在fast+ GDS 标签时为真。

该脚本还内建了单进程分布式初始化（deepspeed_save_model.py），通过设置MASTER_ADDR=localhost、WORLD_SIZE=1等环境变量并以 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.py）：

h = AsyncIOBuilder().load().aio_handle( block_size=8 * (1024**2), # 8 MB 块 queue_depth=8, # 异步队列深度 single_submit=False, # 关闭单提交模式 overlap_events=False, # 事件重叠开关 intra_op_parallelism=1) # 操作内并行度

脚本启动时load_io_ops会先做兼容性探测（torch_save_utils.py）：AsyncIO 与 GDS 扩展只有在构建器is_compatible()通过时才加载。

钉页缓冲（bounce buffer）。CPU 路径的缓冲是一个钉页（pinned memory）张量，默认 64 MB（PINNED_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_handle=h, # AIO 或 GDS 句柄 pinned_tensor=pinned_memory, # CPU 钉页或 GPU 钉页缓冲 double_buffer=not args.single_io_buffer, # 默认双缓冲 num_parallel_writers=1, writer_rank=0) ds_fast_writer = FastFileWriter(file_path=file, config=fast_writer_config) torch.save(f=ds_fast_writer, obj=buffer, ...) ds_fast_writer.close() # 强制 flush 到存储

close()之后的_dump_state()（默认开启，--no-statistics可关闭）会打印底层 I/O 统计。双缓冲（io_buffer_double=True/ 不传--single_io_buffer）让 DMA 填充与异步写出两个阶段重叠，是吞吐最大化的重要开关。

vanilla 对照。test_save直接调用torch.save(f=file, ...)，MockFileWriter与PyFileWriter则是 DeepSpeed 提供的另外两种torch.save()集成写入器，用于在"未启用补丁/未加载 AIO"时给出基准线。

八、加载验证：torch_save_load_model.py 的闭环测试

目录中还有一个不在 README 编号之列但值得关注的辅助脚本 torch_save_load_model.py：它在保存之后立即用torch.load(..., weights_only=True)（或 safetensors）把权重读回，并用torch.allclose(pre_logits, post_logits, atol=1e-3, rtol=1e-3)断言加载前后模型前向输出一致（torch_save_load_model.py）。该脚本支持--regular_torch_save（vanilla 对照）与--safetensors（safetensors 格式）两种旁路，默认走_test_ds_fast_save的 FastPersist 路径。它可以作为"FastPersist 落盘数据与原生 torch 保存等价"的正确性闭环验证手段。

九、复现建议与运行前提

综合 README 与源码，复现 FastPersist 收益的完整路径为：

1. 确认硬件：至少一块 NVMe SSD（推荐多盘 RAID-0 以获得 README 示例量级的带宽）；GDS 路径需要 NVIDIA GPU 与支持 GDS 的驱动/存储环境；
2. 确认软件：torch 2.6.0、DeepSpeed（提供AsyncIOBuilder、GDSBuilder、FastFileWriter等组件）、transformers（requirements.txt 声明）；脚本依赖的 DeepSpeed 扩展需能正常加载；
3. 打补丁：用 serialization_fast_v2.6.0.py 覆盖 torch 安装目录下的torch/serialization.py（先备份原版）；
4. 跑基准：按上文命令依次执行三个脚本，--folder指向 NVMe 挂载点，测模型时--gpu以获得含 GDS 的完整结果；
5. 读日志：在输出中搜索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 版本前需要重新评估补丁兼容性。

• 示例工程