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

本文介绍了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扫描下方二维码免费领取🆓

在这里插入图片描述

在这里插入图片描述

第一阶段(10天):初阶应用

该阶段让大家对大模型 AI有一个最前沿的认识,对大模型 AI 的理解超过 95% 的人,可以在相关讨论时发表高级、不跟风、又接地气的见解,别人只会和 AI 聊天,而你能调教 AI,并能用代码将大模型和业务衔接。

  • 大模型 AI 能干什么?
  • 大模型是怎样获得「智能」的?
  • 用好 AI 的核心心法
  • 大模型应用业务架构
  • 大模型应用技术架构
  • 代码示例:向 GPT-3.5 灌入新知识
  • 提示工程的意义和核心思想
  • Prompt 典型构成
  • 指令调优方法论
  • 思维链和思维树
  • Prompt 攻击和防范

第二阶段(30天):高阶应用

该阶段我们正式进入大模型 AI 进阶实战学习,学会构造私有知识库,扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架,抓住最新的技术进展,适合 Python 和 JavaScript 程序员。

  • 为什么要做 RAG
  • 搭建一个简单的 ChatPDF
  • 检索的基础概念
  • 什么是向量表示(Embeddings)
  • 向量数据库与向量检索
  • 基于向量检索的 RAG
  • 搭建 RAG 系统的扩展知识
  • 混合检索与 RAG-Fusion 简介
  • 向量模型本地部署

第三阶段(30天):模型训练

恭喜你,如果学到这里,你基本可以找到一份大模型 AI相关的工作,自己也能训练 GPT 了!通过微调,训练自己的垂直大模型,能独立训练开源多模态大模型,掌握更多技术方案。

到此为止,大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗?

  • 为什么要做 RAG
  • 什么是模型
  • 什么是模型训练
  • 求解器 & 损失函数简介
  • 小实验2:手写一个简单的神经网络并训练它
  • 什么是训练/预训练/微调/轻量化微调
  • Transformer结构简介
  • 轻量化微调
  • 实验数据集的构建

第四阶段(20天):商业闭环

对全球大模型从性能、吞吐量、成本等方面有一定的认知,可以在云端和本地等多种环境下部署大模型,找到适合自己的项目/创业方向,做一名被 AI 武装的产品经理。

  • 硬件选型
  • 带你了解全球大模型
  • 使用国产大模型服务
  • 搭建 OpenAI 代理
  • 热身:基于阿里云 PAI 部署 Stable Diffusion
  • 在本地计算机运行大模型
  • 大模型的私有化部署
  • 基于 vLLM 部署大模型
  • 案例:如何优雅地在阿里云私有部署开源大模型
  • 部署一套开源 LLM 项目
  • 内容安全
  • 互联网信息服务算法备案

学习是一个过程,只要学习就会有挑战。天道酬勤,你越努力,就会成为越优秀的自己。

如果你能在15天内完成所有的任务,那你堪称天才。然而,如果你能完成 60-70% 的内容,你就已经开始具备成为一名大模型 AI 的正确特征了。

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

😝有需要的小伙伴,可以Vx扫描下方二维码免费领取==🆓

在这里插入图片描述

  • 17
    点赞
  • 20
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值