Open-Sora项目训练与微调完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00965/article/details/148360897

Open-Sora项目训练与微调完整指南

Open-Sora Open-Sora：为所有人实现高效视频制作项目地址: https://gitcode.com/gh_mirrors/op/Open-Sora

前言

Open-Sora是一个开源的视频生成模型项目，能够根据文本描述生成高质量的视频内容。本文将详细介绍如何从零开始训练或微调自己的Open-Sora模型，包括环境准备、数据集处理、模型训练与微调等完整流程。

环境准备

在开始训练前，需要安装一些额外的依赖包：

pip install TensorNVMe  # 用于检查点保存，需要cmake支持
pip install pandarallel  # 用于并行处理

TensorNVMe是一个高性能的存储库，能够加速模型检查点的保存过程。pandarallel则提供了数据并行处理能力，可以显著提高大规模数据集的处理效率。

数据集准备

数据集选择与下载

我们以Pexels 45k数据集为例，这是一个经过剪辑和评分筛选的高质量视频数据集。下载数据集到本地：

mkdir datasets
cd datasets
huggingface-cli download --repo-type dataset hpcai-tech/open-sora-pexels-45k --local-dir open-sora-pexels-45k

下载完成后，需要解压并整理数据集：

cd open-sora-pexels-45k
cat tar/pexels_45k.tar.* > pexels_45k.tar
tar -xvf pexels_45k.tar
mv pexels_45k ..

数据集格式要求

数据集需要提供以下关键信息：

视频路径(path)
文本描述(text)
帧数(num_frames)
高度(height)
宽度(width)
宽高比(aspect_ratio)
分辨率(resolution)
帧率(fps)

数据集预处理

使用提供的脚本处理原始数据集：

# 单进程处理
python scripts/cnv/meta.py --input datasets/pexels_45k.csv --output datasets/pexels_45k_nec.csv --num_workers 0

# 多进程处理(推荐)
python scripts/cnv/meta.py --input datasets/pexels_45k.csv --output datasets/pexels_45k_nec.csv --num_workers 64

预处理过程会根据视频的实际属性计算必要的元数据，这对于支持任意宽高比、分辨率和帧数的训练至关重要。

模型训练

基础训练命令

使用以下命令格式启动训练：

torchrun --nproc_per_node 8 scripts/diffusion/train.py [配置文件路径] --dataset.data-path [数据集路径] [覆盖选项]

例如，使用256px分辨率从头开始训练：

torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv

配置详解

所有训练配置位于configs/diffusion/train/目录下，支持以下特性：

配置继承：使用_base_ = ["config_to_inherit"]语法继承其他配置
命令行参数覆盖：命令行参数会覆盖配置文件中的设置

桶配置(bucket_config)

桶配置用于控制不同训练阶段的数据采样策略，格式如下：

bucket_config = {
    "分辨率": {
        最小帧数: (采样概率, 批量大小),
        ...
    },
    ...
}

例如，256px分辨率下不同帧数的采样策略：

"256px": {
    1: (1.0, 45),   # 图像，100%采样，批量45
    33: (1.0, 12),  # ≥33帧视频，100%采样，批量12
    65: (1.0, 6),   # ≥65帧视频，100%采样，批量6
    ...
}

提供的训练配置

项目提供了多种预设配置：

image.py：仅训练图像生成
stage1.py：256px分辨率视频训练
stage2.py：768px分辨率视频训练(使用序列并行)
stage1_i2v.py：256px分辨率图文视频联合训练
stage2_i2v.py：768px分辨率图文视频联合训练
demo.py：小批量调试配置

模型微调

从Open-Sora v2微调

torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv --model.from_pretrained ckpts/Open_Sora_v2.safetensors

从Flux-dev模型微调

torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv --model.from_pretrained ckpts/flux1-dev-fused-rope.safetensors

多GPU训练

使用ColossalAI进行多节点训练：

colossalai run --hostfile hostfiles --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/stage1.py --dataset.data-path datasets/pexels_45k_necessary.csv

其中hostfiles文件包含各节点的IP地址。

训练恢复

完整恢复

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*

仅恢复优化器状态

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* --start-step 0 --start-epoch 0

注意：如果数据集、批量大小或GPU数量发生变化，数据加载器状态将不再有效。

高级技巧

梯度检查点

通过选择性梯度检查点节省内存：

grad_ckpt_setting = (100, 100)  # (双层检查点数, 单层完整梯度数)
model = dict(
    grad_ckpt_setting=grad_ckpt_setting,
)

可进一步将梯度检查点卸载到CPU：

grad_ckpt_buffer_size = 25 * 1024**3  # 25GB缓冲区

异步检查点保存

启用异步IO加速检查点保存：

--async-io True

大数据集处理

对于超大数据集，可将其分片处理：

python scripts/cnv/shard.py /path/to/dataset.parquet

训练时启用内存高效模式：

--dataset.memory_efficient True

模型推理

训练完成后，可使用以下命令进行推理：

torchrun --nproc_per_node 1 --standalone scripts/diffusion/inference.py configs/diffusion/inference/t2i2v_256px.py --save-dir samples --prompt "raining, sea" --model.from_pretrained outputs/your_experiment/epoch*-global_step*