Open-Sora 训练如何用 TensorNVMe 异步保存 checkpoint？

Open-Sora 训练如何用 TensorNVMe 异步保存 checkpoint【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora训练视频模型时checkpoint 体积大、保存耗时长频繁落盘会拉长每次保存的停顿。Open-Sora 的扩散模型训练支持通过 TensorNVMe 异步保存 checkpoint在 docs/train.md 中--async-io True打开后checkpoint 会在 ColossalAI 的支持下异步写盘官方文档对它的说明是 This will save time for checkpoint saving。本文给出从安装依赖、启动训练到确认 checkpoint 落盘的一条完整操作路径适用于按 docs/train.md 流程训练configs/diffusion/train/下各阶段配置如stage1.py、stage2.py的场景。异步保存的机制与默认状态先明确两件事避免以为打开即可用默认是关闭的。基础配置 configs/diffusion/train/image.py 中写了async_io False而stage1.py、stage2.py都通过_base_ [image.py]继承该默认值所以不额外传参时 checkpoint 走的是同步保存路径。打开方式是命令行覆盖。docs/train.md 的 Config 一节说明 command line arguments override the config file例如--lr 1e-5会覆盖配置中的lr同理在训练命令后追加--async-io True即可覆盖配置里的async_io False。打开后训练脚本 scripts/diffusion/train.py 会调用checkpoint_io.save(..., async_iocfg.get(async_io, False))。保存实现位于 opensora/utils/ckpt.py模型与优化器的分片保存传入use_asyncasync_ioEMA 权重在async_io打开时以ema.safetensors异步写出关闭时以torch.save同步写为ema.pt代码从tensornvme.async_file_io导入AsyncFileWriter见 opensora/utils/ckpt.py即 TensorNVMe 是异步写盘的实际依赖代码在模块加载时设置os.environ[TENSORNVME_DEBUG] 1并在每次保存时把TENSORNVME_DEBUG_LOG指向该 checkpoint 目录下的async_file_io.log见 opensora/utils/ckpt.py排查异步写盘问题时可查看该日志文件。第一步安装 TensorNVMe 依赖docs/train.md 的 Installation 一节给出的安装命令是pip install githttps://github.com/hpcaitech/TensorNVMe.git # requires cmake, for checkpoint saving该命令通过 git 源码安装需要本机已具备 cmake 构建环境文档注释明确其用途就是 for checkpoint saving。另外该节还要求安装pandarallel用于数据集并行处理如果还没装可以一并执行pip install pandarallel # for parallel processing第二步带--async-io True启动训练docs/train.md 给出的训练命令格式为torchrun --nproc_per_node 8 scripts/diffusion/train.py [path/to/config] --dataset.data-path [path/to/dataset] [override options]其中[path/to/config]、[path/to/dataset]是文档中的占位符按文档中的实际示例替换。以 stage 1 配置 pexels 数据集为例在文档示例命令后追加--async-io Truetorchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv --async-io True保存频率由配置控制不需要你手动指定步数configs/diffusion/train/stage1.py 中ckpt_every 2000每 2000 个 actual update step 保存一次、keep_n_latest 20只保留最近 20 个 checkpoint其余由代码自动清理且清理时会跳过eval子目录若用stage2.py对应值是ckpt_every 200、keep_n_latest 20见 configs/diffusion/train/stage2.py。如果要在多机训练时同样开启异步保存colossalai run的命令与单机版一致只是在原有参数后追加--async-io True见 docs/train.md 的 Multi-GPU 一节。如何验证 checkpoint 确实异步保存成功文档没有给出独立的异步保存成功判定命令可以按下面几个文档与源码中实际存在的产物逐项核对checkpoint 目录命名。每次保存会生成epoch{epoch}-global_step{actual_update_step}形式的目录见 opensora/utils/ckpt.py位于实验输出目录默认outputs下的实验工作区下。训练日志中对应一条Saved checkpoint at epoch ... step ... global_step ... to save_dir的记录见 scripts/diffusion/train.py。目录内文件。异步模式下 EMA 权重是ema.safetensors同步模式才是ema.pt目录内还有分片保存的model/、optimizer/、lr_scheduler/以及running_states.json。async_file_io.log也在该目录内TensorNVMe 的调试信息写入其中。能被重新加载。最直接的验证是断点续训docs/train.md 的 Resume training 一节用通配符指定 checkpoint 目录加载 optimizer 与 dataloader 状态torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv --load outputs/your_experiment/epoch*-global_step*其中outputs/your_experiment替换为你实际的实验输出目录。注意CheckpointIO.load会断言目录存在且包含running_states.json见 opensora/utils/ckpt.py如果该文件缺失说明这次保存不完整续训会直接报错。边界与注意事项异步保存依赖 ColossalAI 与 TensorNVMe 两者文档说明异步保存 with the support of ColossalAI而 TensorNVMe 是写盘层依赖安装命令注释标明 for checkpoint saving缺少 cmake 时 TensorNVMe 无法安装训练也就无法依赖该能力。不改变保存频率与清理策略--async-io True只改变写入方式ckpt_every、keep_n_latest仍由配置文件决定。dataloader 状态的限制文档原注如果改变了 dataset、batch size 或 GPU 数量checkpoint 里的 dataloader state 将不再有意义需要只加载 optimizer 状态时按 docs/train.md 的做法追加--start-step 0 --start-epoch 0。异步写盘的同步点代码在每次优化器更新前调用booster.checkpoint_io.synchronize()并在下一次save开头先_sync_io()见 opensora/utils/ckpt.py、scripts/diffusion/train.py即上一次异步写入会阻塞到下一次保存/更新前完成而不是在训练结束时才落盘。【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考