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

AI小白龙*

已于 2024-05-05 10:28:37 修改

阅读量865

点赞数 17

文章标签：深度学习神经网络人工智能 transformer pytorch

于 2024-03-02 11:15:00 首次发布

本文链接：https://blog.csdn.net/2301_81940605/article/details/136380608

版权

本文介绍了OpenDiT项目，一个开源工具，专注于提高DiT模型的训练和推理效率，特别在文本到视频和文本到图像生成方面。项目通过优化并行策略、序列并行方法和高效算子，加速了模型性能，为研究者复现Sora效果提供了新路径。

摘要由CSDN通过智能技术生成

作为 2024 开年王炸，Sora 的出现树立了一个全新的追赶目标，每个文生视频的研究者都想在最短的时间内复现 Sora 的效果。

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

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

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

项目地址：github.com/NUS-HPC-AI-…

OpenDiT 方法介绍

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

图源：www.zhihu.com/people/berk…

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 消耗。

图源：www.zhihu.com/people/berk…

简单来说，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）来运行示例：

css
复制代码
<span><span>conda</span> <span>create -n opendit python=3.10 -y</span></span>

安装 ColossalAI：

xml
复制代码
<span><span>git</span> <span>clone https://github.com/hpcaitech/ColossalAI.gitcd ColossalAI</span></span>

安装 OpenDiT：

bash
复制代码
<span><span>git</span> clone https://github.com/oahzxl/OpenDiTcd OpenDiT</span>

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

css
复制代码
<span><span># Install Triton for fused adaln kernel</span></span>

图像生成

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

css
复制代码
<span><span># Use script</span></span>

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

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 模型进行推理，可以运行如下代码，需要将检查点路径替换为你自己训练的模型。

css
复制代码
<span><span># Use script</span></span>

视频生成

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

css
复制代码
<span><span># train with scipt</span></span>

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

css
复制代码
<span><span># Use script</span></span>

DiT 复现结果

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

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

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

css
复制代码
<span><span>torchrun</span> <span>--standalone --nproc_per_node=8 train.py </span></span>

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以点击下方链接免费领取或者V扫描下方二维码免费领取🆓

在这里插入图片描述