Diffusion预训练成本降低6.5倍,微调硬件成本降低7倍!Colossal-AI完整开源方案低成本加速AIGC产业落地

简介: Diffusion预训练成本降低6.5倍,微调硬件成本降低7倍!Colossal-AI完整开源方案低成本加速AIGC产业落地

如何更好、更快和更便宜地实现训练、微调 AIGC 模型,已成为 AIGC 商业化和应用爆发的最大痛点。


Colossal-AI 基于在大模型民主化的专业技术积累,开源完整 Stable Diffusion 预训练和个性化微调方案,预训练时间加速和经济成本降低 6.5 倍,个性化微调硬件成本降低 7 倍!在个人电脑的 RTX 2070/3050 上即可快速完成微调任务流程,让 Stable Diffusion 等 AIGC 模型的触手可及。


开源地址:

https://github.com/hpcaitech/ColossalAI


火爆的 AIGC 赛道与高昂成本


AIGC(AI-Generated Content 人工智能生成内容)是当前 AI 领域最热门的话题之一,尤其是伴随着 Stable Diffusion、Midjourney、NovelAI、DALL-E 等为代表的文本生成图像的跨模态应用涌现,AIGC 更是火爆出圈,广受关注。


Stable Diffusion 生成图像


由于 AIGC 激发了大量行业需求,它已被视为下一波 AI 浪潮的重要方向之一,业界广泛期望出现基于 AIGC 在文本、音频、图像视频、游戏、元宇宙等技术场景的新技术革命和杀手级应用。AIGC 在相关场景的成功商业化落地,潜在的数万亿美元市场,更是让相关初创公司成为资本宠儿,如 Stability AI、Jasper 等成立仅一两年便已获得上亿美元融资,晋升独角兽行列。


AI 模型规模与性能的同步增长


但高昂的硬件需求和训练成本仍严重阻碍着 AIGC 行业的快速发展。AIGC 应用的出色表现通常建立在 GPT-3 或 Stable Diffusion 等大模型之上,并针对特定下游任务和应用进行微调。以大火的 Stable Diffusion 为例,尽管其背后的 Stability AI 成立不久,却维护了超过 4000 个英伟达 A100 的 GPU 集群,并已为此支出超过 5000 万美元的运营成本,仅 Stable Diffusion v1 版本的模型单次训练便需要 150000 个 A100 GPU Hour。


Diffusion model


Diffusion model(扩散模型)的想法最早在 2015 年的论文 Deep Unsupervised Learning using Nonequilibrium Thermodynamics 被提出,2020  的论文 Denoising Diffusion Probabilistic Models (DDPM)将其推到了一个新的高度,之后基于扩散模型的 DALL-E 2, Imagen, Stable Diffusion 在生成任务上取得了远超生成对抗网络(GAN)、变微分自动编码器(VAE)、自回归模型(AR)等传统生成模型的效果。


扩散模型包含两个过程:前向扩散过程和反向生成过程,前向扩散过程是对一张图像逐渐添加高斯噪声直至变成随机噪音,而反向生成过程是去噪音过程,将一个随机噪音使用多个 U-Net 进行逐渐去噪音直至生成一张图像,这也是扩散模型训练的部分。


Latent Diffusion model


对比传统端到端的深度学习模型,扩散模型的训练过程无疑更为复杂,以 Stable Diffusion 为例,除了扩散模型本身,还有一个 Frozen CLIP Textcoder 来输入 text prompts,以及一个 Autoencoder 实现将高分辨率图像压缩到潜在空间(Latent Space),并在每个 time step 计算 loss。这对训练方案的显存开销,计算速度都提出了更大的挑战。


更低成本——预训练加速与少资源微调


预训练优化


对于预训练而言,一般 batch size 越大,训练速度也越快,Diffusion model 也是类似的。Colossal- AI 通过 ZeRO,Gemini,  Chunk-based 内存管理等策略以及 Flash Attention 模块优化 Cross-attention 计算,极大地降低了 Diffusion model 的训练的显存开销,使用户在 10G 显存的消费级显卡(如 RTX3080)上就可以训练 Diffusion model,在 A100 这样的专用显卡上最大可以直接支持单卡 Batch Size 256 的训练, 对比 stable-diffusion-v1-1 的 FP32 的 DistributedDataParallel (DDP) 训练可以提速 6.5 倍。这意味着数百万美元的训练成本可降低 6.5 倍,极大降低 AIGC 行业训练成本和入场门槛!
Acceleration of Colossal-AI to Stable Diffusion


个性化微调优化


由于 Stable Diffusion 的预训练采用的 LAION-5B 数据集共 5850 亿个图片文本对,需要 240TB 储存空间,再结合模型的复杂性,显然完整预训练的成本极高:Stable Diffusion 的 Stability 团队花费超过 5000 万美元部署了 4,000 块 A100 GPU。对于大多数 AIGC 玩家而言,更切实的选择是使用开源的预训练模型权重来进行微调个性化下游任务。


但其他现有的开源 finetune 方案中使用的训练并行方式主要为 DDP,这导致训练过程中显存占用极大,即使微调也需要至少使用 RTX 3090 或 4090 最高端的消费级显卡才能启动同时,现阶段开源的很多训练框架并没有给出完整的训练配置与脚本,需要用户花费额外时间进行烦琐的补全和调试。


不同于其他解决方案,Colossal-AI 是首个同时开源完整的训练配置参数和训练脚本的方案,让用户可以随时训练出针对新下游任务的最新版细分模型,使用更加灵活且应用范围更广。而且由于 Colossal-AI 引入显存优化等技术,仅在普通个人电脑的单张消费级显卡上(如 GeForce RTX 2070/3050 8GB),即可快速完成微调任务流程,相比 RTX 3090 或 4090 可降低约 7 倍硬件成本,大大降低了使用 Stable Diffusion 等 AIGC 模型的门槛和成本,使用户不再局限于现有的权重推理,方便快捷完成个性化定制服务。对于速度不敏感的任务,还可以进一步使用 Colossal-AI NVMe,即利用低成本的硬盘空间降低显存消耗。


Memory Reduction of Colossal-AI to Stable Diffusion

背后优化技术


ZeRO + Gemini


Colossal-AI 支持使用零冗余优化器 (ZeRO) 的方法来消除内存冗余,与经典的数据并行性策略相比,可极大提高内存使用效率,同时不牺牲计算粒度和通信效率。


Colossal-AI 引入了 Chunk 机制,我们可以进一步提升 ZeRO 的性能。运算顺序上连续的一组参数存入一个 Chunk 中(Chunk 即一段连续的内存空间),每个 Chunk 的大小相同。Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU 之间网络带宽的高效利用,减小了通信次数,同时避免潜在的内存碎片。


Chunk 机制


此外,Colossal-AI 的异构内存空间管理器 Gemini 支持将优化器状态从 GPU 卸载到 CPU ,以节省 GPU 内存占用。可以同时利用 GPU 内存、CPU 内存(由 CPU DRAM 或 NVMe SSD 内存组成)来突破单 GPU 内存墙的限制,进一步扩展了可训练模型规模。


通过 ZeRO + Gemini 提升硬件的模型容量


Flash Attention


LDM(Latent Diffusion Models) 通过在模型架构中引入 cross-attention(交叉注意力层) 来实现多模态训练,使得 Diffusion model 可以更灵活地实现对 class-condition, text-to-image, layout-to-image 的支持。然而 cross-attention 层对比原始 Diffusion model 的 CNN 层增加了额外的计算开销,极大增加了训练成本。



Colossal-AI 通过引入 Flash attention 机制,成功将 attention 的速度提升 104%,将端到端训练的峰值显存减少 23%。Flash attention 是针对长序列 attention 的加速版本,使用 Flatten 来减少 GPU 高带宽内存(HBM)之间的内存读 / 写次数, Flash attention 同时针对块状稀疏的 attention,设计了一个近似的注意力算法,比任何现有的近似 attention 方法都要快。


其他优化


Colossal-AI 还集成了 FP16、activation checkpoint 等常见优化技术。例如,activate checkpoint 通过用计算换取内存来工作。它避免存储整个计算图的所有中间激活用于反向计算,在检查点部分不保存中间激活,而是在反向传递中重新计算它们,进一步降低了显存。而 FP16 在基本不影响精度前提下,将原本的 32 位浮点数运算转为 16 位,降低显存使用,提升计算效率。


快速上手使用


不同于常见的 PyTorch 开源项目,当前火热的 stable diffusion 是基于 PyTorch Lightning 搭建的。PyTorch Lightning 为流行的深度学习框架 PyTorch 提供了简洁易用、灵活高效的高级接口,为广大 AI 研究人员提供了简洁易用的高层次抽象,从而使深度学习实验更易于阅读和再现,已在 GitHub 上收获了 20.5k 颗 Star。


受 PyTorch Lightning 的邀请,Colossal-AI 已集成作为 PyTorch Lightning 的官方大模型解决方案。得益于两者的强强联合,现在 AI 研究者们可以更加高效地训练和使用 diffusion 模型。以训练 stable diffusion model 为例,仅需少量代码即可快捷启动。



from colossalai.nn.optimizer import HybridAdamfrom lightning.pytorch import trainer
class MyDiffuser(LightningModule):    ...
    def configure_sharded_model(self) -> None:        # create your model here        self.model = construct_diffuser_model(...)        ...
    def configure_optimizers(self):        # use the specified optimizer        optimizer = HybridAdam(self.model.parameters(), self.lr)        ...
model = MyDiffuser()trainer = Trainer(accelerator="gpu", devices=1, precision=16, strategy="colossalai")trainer.fit(model)


Colossal-AI 和 PyTorch Lightning 也对 OPT、HuggingFace 等热门模型和社区提供了良好支持及优化。


低成本微调


Colossal-AI 为了满足用户通过较少资源短时间训练出可以生成有自己风格的模型的需求,提供了基于 HuggingFace 上开源的 Stable Diffusion 模型权重进行微调的功能。用户只需简单修改 Dataloader 载入自己的微调数据集并读取预训练权重,简单修改参数配置 yaml 文件并运行训练脚本,便可在个人电脑微调属于自己的个性化模型。


model:
  target: ldm.models.diffusion.ddpm.LatentDiffusion
  params:
    your_sub_module_config:
      target: your.model.import.path
      params:
        from_pretrained: 'your_file_path/unet/diffusion_pytorch_model.bin'
        ...
lightning:
  trainer:
    strategy:
      target: pytorch_lightning.strategies.ColossalAIStrategy
      params:
        ...
 python main.py --logdir /your_log_dir -t -b config/train_colossalai.yaml


快速推理


Colossal-AI 同时支持原生 Stable Diffusion 推理管道,在完成训练或精调后只需直接调用 diffuser 库并加载自己保存的模型参数即可直接进行推理,无需进行其他改动,方便新用户熟悉推理流程并可以让习惯使用原版框架的用户快速上手。



from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
    "your_ColoDiffusion_checkpoint_path"
    ).to("cuda")
image = pipe('your prompt', num_inference_steps=50)["sample"][0]
image.save('file path')



上述推理流程的生成作品


One More Thing


上述针对 Diffusion 为代表的 AIGC 训练优化突破基于面向大模型时代的通用深度学习系统 Colossal-AI,它通过高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等实现高效快速部署 AI 大模型训练和推理,降低 AI 大模型应用成本。自开源以来,Colossal-AI 已经多次在 GitHub 及 Papers With Code 热榜位列世界第一,与众多已有数万 star 的明星开源项目一起受到海内外关注!经国际专家的严格评审,Colossal-AI 已成功入选为 SC、AAAI、PPoPP 等国际 AI 与 HPC 顶级会议的官方教程。


Colossal-AI 应用:更好的蛋白质结构预测解决方案


Colossal-AI 相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用,广受好评。例如,针对生物医药行业的蛋白质结构预测模型 AlphaFold,基于 Colossal-AI 的优化方案 FastFold 成功将单张 GPU 可推理的最大氨基酸序列长度突破至一万,覆盖了 99.9999% 的蛋白质,仅用笔记本电脑上的消费级显卡即可解析 90% 蛋白质。还能进一步对训练、推理进行全流程并行加速,已助力多家新型药物研发企业缩短开发流程,降低研发成本。


开源地址:

https://github.com/hpcaitech/ColossalAI


参考链接

https://github.com/CompVis/stable-diffusion

https://arxiv.org/abs/2205.14135

https://arxiv.org/abs/2112.10752

https://openai.com/blog/triton/

https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
1天前
|
人工智能 Serverless
AI助理精准匹配,为您推荐方案——如何添加一个Stable Difussion图像生成应用
介绍了一种利用AI助手快速获取并搭建Stable Diffusion图像生成应用的方法。用户只需在阿里云官网向AI助手提出需求,即可获得详细的实施方案。随后,按照AI助手提供的方案,通过函数计算部署应用,并进行测试。此过程显著提升了开发效率。
25 1
AI助理精准匹配,为您推荐方案——如何添加一个Stable Difussion图像生成应用
|
1天前
|
人工智能 自然语言处理 搜索推荐
【云栖实录】大模型驱动,开源融合的AI搜索产品发布
本文介绍了2024云栖大会上阿里云发布的产品详情。
|
20天前
|
人工智能 物联网 开发者
魔搭上线AIGC专区,为开发者提供一站式AI创作开发平台
魔搭上线AIGC专区,首批上架157个风格化大模型,专业文生图全免费~
|
18天前
RTX3090可跑,360AI团队开源最新视频模型FancyVideo
【9月更文挑战第23天】近年来,人工智能技术的迅猛发展推动了视频生成领域的进步,但在合成动态、连贯且时长较长的视频方面仍面临挑战。为此,360AI团队提出了FancyVideo模型,通过跨帧文本指导实现更为连贯的视频生成。FancyVideo引入了跨帧文本指导模块(CTGM),包含时间信息注入器(TII)、时间相关性优化器(TAR)和时间特征增强器(TFB)三个组件,分别负责注入帧特定信息、优化相关性和增强时间一致性。这些机制使模型能生成具有连贯动作和丰富运动的视频,适用于动画制作和视频编辑等领域。然而,FancyVideo也存在计算复杂度高和细节真实感提升空间等局限。
32 3
|
2月前
|
人工智能 文字识别 算法
打造全场景、跨领域、多模态的AI工作流 | 开源图像标注工具 X-AnyLabeling v2.4.0 正式发布!
X-AnyLabeling是一款强大的辅助标注工具,集成了AI推理引擎和丰富功能,为图像数据工程师提供一站式解决方案。它支持图像和视频文件的自动标注,提供了包括矩形框、多边形在内的七种标注样式,适应多样化的训练场景需求。X-AnyLabeling内置了多种SOTA级AI模型,如YOLO、SAM系列等,并支持GPU加速和多种数据集格式的导入导出,确保高效的数据处理。此外,它还具备良好的跨平台兼容性,可在多种操作系统上运行,并提供详尽的帮助文档和社区支持,帮助用户轻松上手并解决使用过程中遇到的问题。
120 2
打造全场景、跨领域、多模态的AI工作流 | 开源图像标注工具 X-AnyLabeling v2.4.0 正式发布!
|
2月前
|
存储 人工智能 自然语言处理
无缝融入,即刻智能[二]:Dify-LLM平台(聊天智能助手、AI工作流)快速使用指南,42K+星标见证专属智能方案
【8月更文挑战第8天】无缝融入,即刻智能[二]:Dify-LLM平台(聊天智能助手、AI工作流)快速使用指南,42K+星标见证专属智能方案
无缝融入,即刻智能[二]:Dify-LLM平台(聊天智能助手、AI工作流)快速使用指南,42K+星标见证专属智能方案
|
1月前
|
人工智能 自然语言处理 数据挖掘
【通义】AI视界|性能超越GPT-4o?最强大的开源AI模型来了……
本文介绍了五项最新AI技术动态,包括性能超越GPT-4o的开源AI模型Reflection70B、智谱清言App限时免费的视频通话功能、哈佛医学院研发的癌症诊断AI模型CHIEF、Replit推出的AI编程助手,以及英特尔与日本AIST合作设立的芯片制造研发中心。这些进展展示了AI领域的快速创新与广泛应用。更多详情,请访问通义官网体验。
|
14天前
|
人工智能 自然语言处理 Linux
Llama 3.2:开源可定制视觉模型,引领边缘AI革命
Llama 3.2 系列 11B 和 90B 视觉LLM,支持图像理解,例如文档级理解(包括图表和图形)、图像字幕以及视觉基础任务(例如基于自然语言描述在图像中精确定位对象)。
|
2月前
|
人工智能 安全 Anolis
专访英特尔:开源与 AI 迅猛发展,龙蜥社区提供高效合作平台
龙蜥社区为开源技术、AI 技术提供了一个非常好的高效合作的平台。

热门文章

最新文章