火爆全球的 Stable Diffusion，预训练和微调成本，正式被打下来了！
现在，预训练的速度被提升到 6.5 倍，成本直降 85%；
微调只需 RTX 2070/3050 即可搞定，硬件成本直降近九成。
这就是国产项目 Colossal-AI 最新开源的加速方案。
今年，Stable Diffusion 可谓火爆圈内圈外，而且最近凭借免费开源还效果好的特点，进一步出圈，让人看到了 AIGC 的巨大潜力和前景。
艺术圈和 AI 圈大神都争相拿它创作，GitHub star 数已达 33.5k，背后公司估值甚至飙升至 10 亿美金。
Stable Diffusion 由此也成为扩散模型中最火出圈的项目。
BUT，它唯一的缺点是，训练成本真的很高！
仅 Stable Diffusion v1 版本的模型单次训练便需要 150000 个 A100 GPU Hour；
微调至少需要 RTX 3090/4090 这类高端消费级显卡。
其背后公司 Stability AI 为了维护运营和算力需求，拥有一个由 4000 块 A100 组成的 GPU 集群，光成本就超过 5000 万美元（折合人民币约 3.6 亿）。
实际上，由于扩散模型本身原理机制特殊，导致它本身对硬件显存、计算速度的要求都很高，而且使用的数据库也是千亿量级的。
所以想要提升训练速度、降低成本，还真不是件容易事。
那么，Colossal-AI 到底是如何实现的？
从显存入手
首先来解释，为什么扩散模型对硬件的要求高。
扩散模型主要包括两个过程：前向扩散和反向生成。
前向扩散是对一张图像逐渐添加高斯噪声，直至变成随机噪声。
反向生成是去噪声过程，将一个随机噪声使用多个 U-Net 进行逐渐去噪，直至生成一张图像，这也是扩散模型的训练过程。
相比于传统端到端深度学习模型，扩散模型的训练过程显然更为复杂。
其中 Stable Diffusion 除了扩散模型本身，还包含一个 Frozen CLIP Textcoder 来输入文本提示词（text prompts）、一个自动编码器（Autoencoder）将高分辨率图像压缩到潜在空间（Latent Space），并且还要在每个 time step 计算 loss。
这样一来，训练方案对显存开销、计算速度都提出了极大挑战。
针对如上问题，Colossal-AI 分别从预训练、微调两个步骤上进行了优化。
预训练阶段，主要做的就是降低显存开销。
一般来说 batch size 越大，训练速度越快，扩散模型亦是如此。
通过 ZeRO、Gemini、Chunk-based 内存管理等策略以及 Flash Attention 模块优化 Cross-attention 计算，Colossal-AI 可以让用户在 8GB 显存的普通消费级显卡（如 RTX 2070/3050）上训练扩散模型。
如果在 A100 上，则能将单卡 Batch Size 直接扩大到 256。
对比 Stable-diffusion-v1 FP32 的 Distributed Data Parallel (DDP) ，训练可以提速 6.5 倍。
由此可以大幅缩短训练时长，从而降低高达数百万美元的训练成本。
在微调上，对于大多数 AIGC 玩家而言，都倾向于选择使用开源的预训练模型权重来进行微调个性化下游任务。
一方面是由于扩散模型本身复杂，另一方面是 Stable Diffusion 预训练采用的是 LAION-5B 数据集，包含 5850 亿个图片文本对，需要 240TB 储存空间。
但现有的很多开源微调方案中，使用的训练并行方式主要为 DDP，这导致训练过程中占用的显存很多。
即使微调也至少需要 RTX 3090/4090 这类最高端的消费级显卡。
加之，现阶段很多开源的训练框架都没有给出完整的训练配置与脚本，还需要自己来做一些补全和调试工作。
基于现状，Colossal-AI 提出了首个同时开源完整训练配置参数和训练脚本的方案。
这种方案一方面可以让用户不再局限于现有的权重推理，更加灵活地训练出针对新下游任务的最新版细分模型，扩大应用范围；
另一方面，由于引入了显存优化等技术，方案支持在普通个人电脑的单张显卡（如 GeForce RTX 2070/3050 8GB）上快速完成微调任务。
要知道，RTX 2070 现在的价格可比 3090/4090 便宜了将近 9 成。
此外，对于速度不敏感的任务，还可以进一步使用 Colossal-AI NVMe，即利用低成本的硬盘空间降低显存消耗。
优化秘籍是啥？
所以显存开销是通过哪些技术降下去的？
主要基于三个方面：
ZeRO+Gemini
Flash Attention
FP16、acitvation checkpoint 等常见优化技术
先来看 ZeRO+Gemini。
ZeRO 是零冗余优化器。
这种方法主要通过切分优化器状态、梯度、模型参数，使 GPU 仅保存当前计算所需的部分，从而来消除数据并行、模型并行中存在的内存冗余。
和经典的数据并行策略相比，它能极大提高内存使用效率，同时不牺牲计算粒度和通信效率。
与此同时，Colossal-AI 还引入了 Chunk 机制进一步提升 ZeRO 性能。
将运算顺序上连续的一组参数存入一个 Chunk 中（Chunk 即一段连续的内存空间），每个 Chunk 的大小相同。
Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU 之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。
Gemini 是 Colossal-AI 的异构内存空间管理器。
它支持将优化器从 GPU 卸载到 CPU，以节省 GPU 内存占用。
打破 GPU 内存墙的限制，可以实现同时利用 GPU 内存和 CPU 内存（由 CPU DRAM 或 NVMe SSD 内存组成），从而进一步扩展可训练模型规模。
第二方面是 Flash Attention 机制。
LDM（Latent Diffusion Models）是改善扩散模型效率的一种方法，它通过在模型框架中引入交叉注意力机制（Cross-Attention）来实现多模态训练，使得扩散模型可以更灵活地实现对 class-condition, text-to-image, layout-to-image 的支持。
但是 Cross-Attention 层对比原始扩散模型的 CNN 层增加了额外的计算开销，极大增加了训练成本。
相比之下，Flash Attention 是针对长序列 Attention 的加速版本。
它使用 Flatten 来减少 GPU 高带宽内存（HBM）之间的内存读 / 写次数，Flash Attention 同时针对块状稀疏的 Attention，设计了一个近似的注意力算法，比现有任何 attention 方法都要快。
它能将 Attention 的速度提升 104%，将端到端训练的峰值显存减少 23%。
除了以上两方面，Colossal-AI 还集成了 FP16、activation checkpoint 等常见优化技术。
例如，activate checkpoint 通过用计算换取内存来工作。
它避免存储整个计算图的所有中间激活用于反向计算，在检查点部分不保存中间激活，而是在反向传递中重新计算它们，进一步降低了显存。
而 FP16 在基本不影响精度前提下，将原本的 32 位浮点数运算转为 16 位，降低显存使用，提升计算效率。
上手也超简单
到了实操环节，Colossal-AI 只需少量代码，即可快捷启动扩散模型训练。
不同于常见 PyTorch 开源项目，Stable Diffusion 是基于 PyTorch Lighting 搭建的。
PyTorch Lightning 为流行的深度学习框架 PyTorch 提供了简洁易用、灵活高效的高级接口，为广大 AI 研究人员提供了简洁易用的高层次抽象，从而使深度学习实验更易于阅读和再现，
受 PyTorch Lightning 的邀请，Colossal-AI 已集成作为 PyTorch Lightning 的官方大模型解决方案。
以 Stable Diffusion 为例，启动代码如下：
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
实际上，如上提到的 Colossal-AI 不只能用来加速训练扩散模型，它面向各种大模型。
通过多维自动并行、异构内存管理、大规模优化库、自适应任务调度等方法，Colossal-AI 可以高效快速部署 AI 大模型训练和推理，降低应用成本。
自开源以来，Colossal-AI 多次在 GitHub 及 Paper With Code 热榜上位列世界第一。
经国际专家的严格评审，Colossal-AI 已成功入选为 SC、AAAI、PPoPP 等国际 AI 与 HPC 顶级会议的官方教程。
目前，其相关解决方案已经成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用。
例如针对生物医药行业的蛋白质结构预测模型 AlphaFold，基于 Colossal-AI 的优化方案，FastFold 成功将单张 GPU 可推理氨基酸序列长度扩增到一万，覆盖 99.9999% 的蛋白质。
只用笔记本上的消费级显卡就能解析 90% 蛋白质。
而且还能对训练、推理进行全流程并行加速，目前已助力多家新型药物研发企业缩短开发流程，降低开发成本。
传送门
开源地址：
https://github.com/hpcaitech/ColossalAI
参考链接：
[1]https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b
[2]https://github.com/CompVis/stable-diffusion
[3]https://arxiv.org/abs/2205.14135
[4]https://arxiv.org/abs/2112.10752
[5]https://openai.com/blog/triton/
本文来自微信公众号：量子位 （ID：QbitAI），作者：明敏
广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。