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行业训练成本和入场门槛！

 

• 个性化微调优化

由于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，即利用低成本的硬盘空间降低显存消耗。

背后优化技术

• 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 HybridAdam
from 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/