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 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/

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

© THE END

转载请联系本公众号获得授权

投稿或寻求报道：content@jiqizhixin.com