Open-Sora 训练如何用梯度检查点与 CPU 卸载节省显存？

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 做检查点）和把检查点激活卸载到 CPU（grad_ckpt_buffer_size）。这篇文章围绕"在启动训练前把显存压下来"这一任务，说明这两个开关分别写在哪里、取什么值、如何用真实命令启动，以及如何在训练日志里确认显存确实降了。

适用对象：在configs/diffusion/train/配置基础上训练或微调 Open-Sora 的开发者。训练入口是 scripts/diffusion/train.py，配置基类是 configs/diffusion/train/image.py。

三个控制显存的配置项分别在哪里

训练脚本按下面三个键读取配置，三者作用不同，容易混淆：

在 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_setting=grad_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.py：

grad_ckpt_settings = (8, 100) ... model = dict(grad_ckpt_settings=grad_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, dtype=dtype, pin_memory=True)，即按训练精度（配置默认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 git+https://github.com/hpcaitech/TensorNVMe.git # requires cmake, for checkpoint saving pip install pandarallel # for parallel processing

环境本身见 README 的 Installation（Python 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、L181）
• log_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 GPU（140GB 显存）上搜出来的（docs/train.md）。也就是说，这些配置假设了较大的单卡显存；在显存更小的卡上，应先用demo.py（batch 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