Open-Sora 训练如何用梯度检查点与 CPU 卸载节省显存【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora当你用 Open-Sora 的configs/diffusion/train/配置如stage1.py、stage2.py启动扩散模型训练时显存往往是最先吃紧的资源。仓库文档在 docs/train.md 的 Gradient Checkpointing 一节给出了两个直接作用于显存的开关选择性梯度检查点grad_ckpt_settings控制对多少 Transformer block 做检查点和把检查点激活卸载到 CPUgrad_ckpt_buffer_size。这篇文章围绕在启动训练前把显存压下来这一任务说明这两个开关分别写在哪里、取什么值、如何用真实命令启动以及如何在训练日志里确认显存确实降了。适用对象在configs/diffusion/train/配置基础上训练或微调 Open-Sora 的开发者。训练入口是 scripts/diffusion/train.py配置基类是 configs/diffusion/train/image.py。三个控制显存的配置项分别在哪里训练脚本按下面三个键读取配置三者作用不同容易混淆配置键位置作用现有配置的真实取值grad_checkpoint顶层梯度检查点总开关train.py据此调用set_grad_checkpointimage.py中为Truegrad_ckpt_settingsmodel dict(...)内部选择性检查点的元组(double, single)stage1.py、image.py为(8, 100)stage2.py、high_compression.py、stage2_i2v.py为(100, 100)grad_ckpt_buffer_size顶层CPU 卸载缓冲区大小元素数仅stage2_i2v.py设置为25 * 1024**3在 scripts/diffusion/train.py 中grad_ckpt_buffer_size通过cfg.get(grad_ckpt_buffer_size, 0)读取若大于 0 才调用GLOBAL_ACTIVATION_MANAGER.setup_buffer(...)预分配一块 pinned CPU 缓冲区grad_checkpoint通过cfg.get(grad_checkpoint, True)读取为真时对模型执行set_grad_checkpoint(model)。也就是说文档示例把grad_ckpt_buffer_size写成顶层变量与实际代码一致而元组必须挂在model字典里。如何配置选择性梯度检查点文档给出的做法是直接改配置文件。docs/train.md 的描述是We support selective gradient checkpointing to save memory. Thegrad_ckpt_settingis a tuple, the first element is the number of dual layers to apply gradient checkpointing, the second element is the number of single layers to apply full gradient. A very large number will apply full gradient to all layers.文档示例写的是grad_ckpt_setting (100, 100) model dict( grad_ckpt_settinggrad_ckpt_setting, )这里有一个必须注意的命名坑文档段落里写的是单数grad_ckpt_setting但仓库里所有真实配置与模型代码用的都是复数grad_ckpt_settings。模型 opensora/models/mmdit/model.py 读取的是self.config.grad_ckpt_settings并按它是否为空决定用哪条前向路径if self.config.grad_ckpt_settings: self.forward self.forward_selective_ckpt else: self.forward self.forward_ckpt如果你照抄文档里的单数名等于写了一个模型不读取的键选择性检查点不会生效。请以真实配置为准例如 configs/diffusion/train/stage1.pygrad_ckpt_settings (8, 100) ... model dict(grad_ckpt_settingsgrad_ckpt_settings)关于元组第二个元素的语义文档与源码存在一处口径差异需如实指出文档说第二个数是对单层块使用完整梯度的数量数值很大则所有层用完整梯度而 forward_selective_ckpt 实际是取double_blocks[:第一个数]与single_blocks[:第二个数]去做auto_grad_checkpoint即前 N 个 block 做检查点其余 block 走普通前向。因此在实际调参时请以配置文件和forward_selective_ckpt的行为为准文档的措辞可作为它想表达的意图参考。若需要最大化的检查点覆盖用(100, 100)这类大于实际 block 数量的值真实配置里stage2.py、high_compression.py即如此。如何用 grad_ckpt_buffer_size 把检查点激活卸载到 CPU文档进一步说 To further save memory, you can offload gradient checkpointing to CPU对应顶层配置grad_ckpt_buffer_size 25 * 1024**3 # 25GB文档注释把它标为 25GB。在代码里scripts/diffusion/train.py 读到该值默认0即不启用后调用setup_buffer(numel, dtype)而 ActivationManager.setup_buffer 会执行torch.empty(numel, dtypedtype, pin_memoryTrue)即按训练精度配置默认dtype bf16见 image.py预分配一块固定的 CPU pinned 缓冲区。真实使用示例是 configs/diffusion/train/stage2_i2v.py它在stage2.py基础上追加了grad_ckpt_buffer_size 25 * 1024**3。这块缓冲区是固定大小的超出会报错ActivationManager.offload 在avail_offset size total_size时抛出RuntimeError: Activation buffer is full。所以这是一个明确的失败信号——如果你启用了 CPU 卸载却看到这个错误说明缓冲区装不下被卸载的激活需要调大grad_ckpt_buffer_size或减少同时驻留在缓冲区里的检查点激活量。如何启动训练准备环境。除主页面安装外docs/train.md 要求额外安装pip install githttps://github.com/hpcaitech/TensorNVMe.git # requires cmake, for checkpoint saving pip install pandarallel # for parallel processing环境本身见 README 的 InstallationPython 3.10、torch 2.4.0、pip install -v .、xformers、flash-attn。数据集需为csv或parquet。用自定义数据时至少要包含这些列path,text,num_frames,height,width,aspect_ratio,resolution,fps启动命令格式docs/train.md 的 Training 一节torchrun --nproc_per_node 8 scripts/diffusion/train.py [path/to/config] --dataset.data-path [path/to/dataset] [override options]以stage1.py元组(8, 100)为例torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv[path/to/config]替换为configs/diffusion/train/下的目标配置[path/to/dataset]替换为你的csv/parquet路径。文档还给出两条配置覆盖规则_base_继承以及命令行覆盖配置——文档以--lr 1e-5覆盖顶层lr、--dataset.data-path ...覆盖dataset字典里的值为例。按同样的顶层键规则可用--grad_ckpt_buffer_size 整数覆盖顶层的grad_ckpt_buffer_size而grad_ckpt_settings元组建议直接改配置文件文档的示例就是改配置。如果你的显存/卡数很紧文档提供了 configs/diffusion/train/demo.py它继承stage1.py并把各 bucket 的 batch size 设为 1明确用于调试a demo config demo.py with small batch size for debugging。它是验证改完配置能否跑通、显存是否降下来的低成本路径可作为可选分支不替代正式的stage1/stage2训练。如何在日志中确认显存降了训练入口 scripts/diffusion/train.py 在构建每个组件后都会打印显存读者据此对比开关前后的数值print_mem(before prepare_dataloader)/print_mem(after prepare_dataloader)L175、L181log_cuda_memory(...)在diffusion、EMA、autoencoder、t5、clip、optimizer、boost、sharding EMA各阶段各记一次L194 至 L344训练结束时log_cuda_max_memory(final)记录峰值显存L650对比方法同一配置分别在仅开grad_checkpoint/默认元组与调低grad_ckpt_settings或追加grad_ckpt_buffer_size下各跑一次比较日志里各阶段的 CUDA 显存与final峰值。若启用 CPU 卸载时出现Activation buffer is full回到上一节调大缓冲区。边界与限制文档明确bucket_config里的 batch size 是在 H200 GPU140GB 显存上搜出来的docs/train.md。也就是说这些配置假设了较大的单卡显存在显存更小的卡上应先用demo.pybatch size 1验证再按bucket_config的(sampling probability, batch size)元组下调各分辨率的 batch size。grad_ckpt_buffer_size默认0即默认不启用 CPU 卸载只有显式设置大于 0 的值才会预分配缓冲区。文档的grad_ckpt_setting单数与代码/配置实际的grad_ckpt_settings复数不一致元组第二个元素的语义在文档与forward_selective_ckpt之间也存在口径差异。调参时以真实配置和模型代码为准不要把文档示例中的单数名直接照抄。下一步确认显存达标后可用文档 Fine-tuning 一节的--model.from_pretrained从已有权重继续训练或用--load断点续训。【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
