明民从凹秘史出来了。

量子比特|公众号QbitAI

火热的全球可持续发展，前期训练和微调费用，正式受到打击！

目前，事前训练速度提高了6.5倍，费用飙升到了85%。

微调可以通过RTX 2070/3050处理，硬件成本接近90%。

这就是国产项目Colossal-AI最新开源的加速方案。

开源地址：

，今年的Stable Diffusion可以说是火爆圈外，最近免费开源以效果好的特点进一步圈出，展示了AIGC的巨大潜力和前景。

艺术圈和AI圈的大神们都争相创作了它，GitHub star的数量达到了33.5k，随后公司价值甚至飙升到了10亿美元。

至此，Stable Diffusion成为分布式模型中最热门的项目。

BUT，它唯一的缺点是训练成本真的很高！

仅Stable Diffusion v1版本的型号一次培训就需要150，000个A100 GPUHOUR。

微调至少需要高端消费级显卡，如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)。

因此，训练方案对视频内存开销和计算速度提出了巨大挑战。

针对上述问题，Colossal-AI分别进行了两个阶段的优化：前期培训和微调。

在训练前阶段，主要做的事情是减少视频内存开销。

一般来说，Batch size越大，训练速度越快，扩散模型也是如此。

Colossal-AI通过ZeRO、Gemini、Chunk-based memory management和Flash Attention module优化了Cross-attention计算，使8GB图形内存适用于消费者的图形卡(

在A100上，单卡Batch Size可以直接扩展到256。

与Stable-diffusion-v1 FP32上的分布式数据并行(DDP)相比，培训速度可提高6.5倍。

因此，训练时间大幅缩短，可以节省数百万美元的教育费用。

0/cda537ee509b4e3794796ad33ef0d15c?_iz=31825&from=ar;x-expires=1706865440&x-signature=qMkXcZUoQdm4v46AgZubxCQ%2B%2Bjo%3D&index=3" width="640" height="305"/>

在微调上，对于大多数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之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。

△Chunk机制

Gemini是Colossal-AI的异构内存空间管理器。

它支持将优化器从GPU卸载到CPU，以节省GPU内存占用。

打破GPU内存墙的限制，可以实现同时利用GPU内存和CPU内存（由 CPU DRAM 或 NVMe SSD内存组成），从而进一步扩展可训练模型规模。

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

第二方面是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%。

△ Flash Attention

除了以上两方面，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为例，启动代码如下(Python)：

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

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

低成本微调

在微调方面，Colossal-AI提供了基于HuggingFace上开源的Stable Diffusion模型权重进行微调的功能。

用户只需简单修改Dataloader载入自己的微调数据集并读取预训练权重，再修改参数配置yaml文件并运行训练脚本，便可在个人电脑微调属于自己的个性化模型。

model:
target: ldm.models.di
params:
your_sub_module_config:
target: your.model.im
params:
from_pretrained: 'your_file_path/une;
...
lightning:
trainer:
strategy:
target:
params:
...
python main.py --logdir /your_log_dir -t -b config

快速推理

Colossal-AI同时支持原生Stable Diffusion推理管道。

在完成训练或精调后，只需直接调用diffuser库并加载自己保存的模型参数，即可直接进行推理，无需进行其他改动。

这样可以方便新用户熟悉推理流程，并让习惯使用原版框架的用户快速上手。

from diffusers import StableDiffusionPipeline
pipe = S(
"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%蛋白质。

而且还能对训练、推理进行全流程并行加速，目前已助力多家新型药物研发企业缩短开发流程，降低开发成本。

传送门

开源地址：

参考链接：
[1]
[2]
[3]
[4]
[5]

— 完 —

量子位 QbitAI · 头条号签约

关注我们，第一时间获知前沿科技动态