作为 2024 开年王炸，Sora 的出现树立了一个全新的追赶目标，每个文生视频的钻研者都想在最短的时间内复现 Sora 的效果。
根据 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 python=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-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./--global-option="--cuda_ext" --global-option="--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
感兴趣的读者可以查看项目主页，了解更多钻研内容。
参考链接：https://www.zhihu.com/people/berkeley-you-yang