AI-GeneratedContent：StableDiffusion引领文本生成图像的跨模态革命

文章主题：AIGC, 文本生成图像, 扩散模型, 人工智能生成内容

点击下方卡片，关注“FightingCV”公众号

回复“AI”即可获得超100G人工智能的教程

点击进入→FightingCV交流群

来源 | 机器之心编辑部

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

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

开源地址：

https://github.com/hpcaitech/ColossalAI

火爆的 AIGC 赛道与高昂成本

AIGC（人工智能生成内容）无疑是当今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 万美元，并配备了 4000 块 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，即利用低成本的硬盘空间降低显存消耗。

在本篇文章中，我们将探讨如何降低Colossal-AI的内存消耗，让它变得更加稳定。首先，我们需要了解Colossal-AI的工作原理以及为什么会出现内存消耗过大的问题。经过分析，我们发现这是由于模型训练过程中大量数据的处理所导致的。因此，为了降低Colossal-AI的内存消耗，我们可以采取以下几种方法：1. **数据筛选与预处理**：在进行模型训练之前，我们对数据进行了筛选和预处理，移除了那些对于训练过程无关紧要的数据，从而减少了训练过程中的内存消耗。2. **模型结构优化**：通过对Colossal-AI的模型结构进行优化，我们采用了更加紧凑的网络结构，降低了模型的复杂度，进一步减少了内存消耗。3. **参数压缩技术**：我们使用了参数压缩技术，如知识蒸馏和网络剪枝等方法，将模型中的参数数量大大减少，从而降低了内存消耗。4. **动态调整学习率**：在训练过程中，我们采用了动态调整学习率的方法，使得模型在学习过程中能够更加高效地利用内存资源。通过以上方法，我们成功地降低了Colossal-AI的内存消耗，使得模型变得更加稳定。同时，这些方法也为其他大规模AI模型的优化提供了借鉴，有助于提高整个AI领域的性能。希望这篇文章能为你在使用Colossal-AI时提供一些帮助！💡

背后优化技术

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(filepath)

上述推理流程的生成作品

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

往期回顾

基础知识

【CV知识点汇总与解析】|损失函数篇

【CV知识点汇总与解析】|激活函数篇