想训练类Sora模型吗？尤洋团队OpenDiT实现80%加速

根据 OpenAI 披露的技术报告，Sora 的核心技术点之一是将视觉数据转化为 patch 的统一表征形式，并通过 Transformer 和扩散模型结合，展现了卓越的扩展（scale）特性。在报告公布后，Sora 核心研发成员 William Peebles 和纽约大学计算机科学助理教授谢赛宁合著的论文《Scalable Diffusion Models with Transformers》就成了众多研究者关注的重点。大家希望能以论文中提出的 DiT 架构为突破口，探索复现 Sora 的可行路径。

最近，新加坡国立大学尤洋团队开源的一个名为 OpenDiT 的项目为训练和部署 DiT 模型打开了新思路。

OpenDiT 是一个易于使用、快速且内存高效的系统，专门用于提高 DiT 应用程序的训练和推理效率，包括文本到视频生成和文本到图像生成。

项目地址：https://github.com/NUS-HPC-AI-Lab/OpenDiT

OpenDiT 方法介绍

OpenDiT 提供由 Colossal-AI 支持的 Diffusion Transformer (DiT) 的高性能实现。在训练时，视频和条件信息分别被输入到相应的编码器中，作为DiT模型的输入。随后，通过扩散方法进行训练和参数更新，最终将更新后的参数同步至EMA（Exponential Moving Average）模型。推理阶段则直接使用EMA模型，将条件信息作为输入，从而生成对应的结果。

图源：https://www.zhihu.com/people/berkeley-you-yang

OpenDiT 利用了 ZeRO 并行策略，将 DiT 模型参数分布到多台机器上，初步降低了显存压力。为了取得更好的性能与精度平衡，OpenDiT 还采用了混合精度的训练策略。具体而言，模型参数和优化器使用 float32 进行存储，以确保更新的准确性。在模型计算的过程中，研究团队为 DiT 模型设计了 float16 和 float32 的混合精度方法，以在维持模型精度的同时加速计算过程。

DiT 模型中使用的 EMA 方法是一种用于平滑模型参数更新的策略，可以有效提高模型的稳定性和泛化能力。但是会额外产生一份参数的拷贝，增加了显存的负担。为了进一步降低这部分显存，研究团队将 EMA 模型分片，并分别存储在不同的 GPU 上。在训练过程中，每个 GPU 只需计算和存储自己负责的部分 EMA 模型参数，并在每次 step 后等待 ZeRO 完成更新后进行同步更新。

FastSeq

在 DiT 等视觉生成模型领域，序列并行性对于有效的长序列训练和低延迟推理是必不可少的。

然而，DeepSpeed-Ulysses、Megatron-LM Sequence Parallelism 等现有方法在应用于此类任务时面临局限性 —— 要么是引入过多的序列通信，要么是在处理小规模序列并行时缺乏效率。

为此，研究团队提出了 FastSeq，一种适用于大序列和小规模并行的新型序列并行。FastSeq 通过为每个 transformer 层仅使用两个通信运算符来最小化序列通信，利用 AllGather 来提高通信效率，并策略性地采用异步 ring 将 AllGather 通信与 qkv 计算重叠，进一步优化性能。

算子优化

在 DiT 模型中引入 adaLN 模块将条件信息融入视觉内容，虽然这一操作对模型的性能提升至关重要，但也带来了大量的逐元素操作，并且在模型中被频繁调用，降低了整体的计算效率。为了解决这个问题，研究团队提出了高效的 Fused adaLN Kernel，将多次操作合并成一次，从而增加了计算效率，并且减少了视觉信息的 I/O 消耗。

图源：https://www.zhihu.com/people/berkeley-you-yang

简单来说，OpenDiT 具有以下性能优势：

1、在 GPU 上加速高达 80%，50%的内存节省

• 设计了高效的算子，包括针对DiT设计的 Fused AdaLN，以及 FlashAttention、Fused Layernorm 和HybridAdam。
• 采用混合并行方法，包括 ZeRO、Gemini 和 DDP。对 ema 模型进行分片也进一步降低了内存成本。

2、FastSeq：一种新颖的序列并行方法

• 专为类似 DiT 的工作负载而设计，在这些应用中，序列通常较长，但参数相比于 LLM 较小。
• 节点内序列并行可节省高达 48% 的通信量。
• 打破单个 GPU 的内存限制，减少整体训练和推理时间。

3、易于使用

• 只需几行代码的修改，即可获得巨大的性能提升。
• 用户无需了解分布式训练的实现方式。

4、文本到图像和文本到视频生成完整 pipeline

• 研究人员和工程师可以轻松使用 OpenDiT pipeline 并将其应用于实际应用，而无需修改并行部分。
• 研究团队通过在 ImageNet 上进行文本到图像训练来验证 OpenDiT 的准确性，并发布了检查点（checkpoint）。

安装与使用

要使用 OpenDiT，首先要安装先决条件：

• Python >= 3.10
• PyTorch >= 1.13（建议使用 >2.0 版本）
• CUDA >= 11.6

建议使用 Anaconda 创建一个新环境（Python >= 3.10）来运行示例：

conda create -n opendit pythnotallow=3.10 -y

conda activate opendit

安装 ColossalAI：

git clone https://github.com/hpcaitech/ColossalAI.gitcd ColossalAI

git checkout adae123df3badfb15d044bd416f0cf29f250bc86

pip install -e .

安装 OpenDiT：

git clone https://github.com/oahzxl/OpenDiTcd OpenDiT

pip install -e .

（可选但推荐）安装库以加快训练和推理速度：

# Install Triton for fused adaln kernel

pip install triton

# Install FlashAttention

pip install flash-attn

# Install apex for fused layernorm kernel

git clone https://github.com/NVIDIA/apex.gitcd apex

git checkout 741bdf50825a97664db08574981962d66436d16a

pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-optinotallow=--cpp_ext" --config-settings "--build-optinotallow=--cuda_ext" ./--global-optinotallow="--cuda_ext" --global-optinotallow="--cpp_ext"

图像生成

你可以通过执行以下命令来训练 DiT 模型：

# Use script

bash train_img.sh# Use command line

torchrun --standalone --nproc_per_node=2 train.py \

    --model DiT-XL/2 \

    --batch_size 2

默认禁用所有加速方法。以下是训练过程中一些关键要素的详细信息：

• plugin: 支持 ColossalAI、zero2 和 ddp 使用的 booster 插件。默认是 zero2，建议启用 zero2。
• mixed_ precision：混合精度训练的数据类型，默认是 fp16。
• grad_checkpoint: 是否启用梯度检查点。这节省了训练过程的内存成本。默认值为 False。建议在内存足够的情况下禁用它。
• enable_modulate_kernel: 是否启用 modulate 内核优化，以加快训练过程。默认值为 False，建议在 GPU < H100 时启用它。
• enable_layernorm_kernel: 是否启用 layernorm 内核优化，以加快训练过程。默认值为 False，建议启用它。
• enable_flashattn: 是否启用 FlashAttention，以加快训练过程。默认值为 False，建议启用。
• sequence_parallel_size：序列并行度大小。当设置值 > 1 时将启用序列并行。默认值为 1，如果内存足够，建议禁用它。

如果你想使用 DiT 模型进行推理，可以运行如下代码，需要将检查点路径替换为你自己训练的模型。

# Use script

bash sample_img.sh# Use command line

python sample.py --model DiT-XL/2 --image_size 256 --ckpt ./model.pt

视频生成

你可以通过执行以下命令来训练视频 DiT 模型：

# train with scipt

bash train_video.sh# train with command line

torchrun --standalone --nproc_per_node=2 train.py \

    --model vDiT-XL/222 \

    --use_video \

    --data_path ./videos/demo.csv \

    --batch_size 1 \

    --num_frames 16 \

    --image_size 256 \

    --frame_interval 3

# preprocess

# our code read video from csv as the demo shows

# we provide a code to transfer ucf101 to csv format

python preprocess.py

使用 DiT 模型执行视频推理的代码如下所示：

# Use script

bash sample_video.sh# Use command line

python sample.py \

    --model vDiT-XL/222 \

    --use_video \

    --ckpt ckpt_path \

    --num_frames 16 \

    --image_size 256 \

    --frame_interval 3

DiT 复现结果

为了验证 OpenDiT 的准确性，研究团队使用 OpenDiT 的 origin 方法对 DiT 进行了训练，在 ImageNet 上从头开始训练模型，在 8xA100 上执行 80k step。以下是经过训练的 DiT 生成的一些结果：

损失也与 DiT 论文中列出的结果一致：

要复现上述结果，需要更改 train_img.py 中的数据集并执行以下命令：

torchrun --standalone --nproc_per_node=8 train.py \

    --model DiT-XL/2 \

    --batch_size 180 \

    --enable_layernorm_kernel \

    --enable_flashattn \

    --mixed_precision fp16

感兴趣的读者可以查看项目主页，了解更多研究内容。