硬件预算最高直降46倍！低成本上手AIGC和千亿大模型，一行代码自动并行，Colossal-AI再升级

作为 当下最火热的开源 AI 大模型解决方案，Colossal-AI 已收获 Github Star 七千多颗 ，此前在 Stable Diffusion、GPT-3、AlphaFold 等大模型上展现卓越性能优势。针对 AI 大模型落地成本高昂这一痛点，Colossal-AI 本次更新聚焦于降低大模型应用成本以及增强易用性，包括：

• Stable Diffusion 2.0 低成本训练 / 微调 / 推理，最多可降低显存消耗 5.6 倍，使硬件成本直降 46 倍，一行代码即可使用；
• 1750 亿参数 BLOOM 模型单机推理，显存消耗降低 4 倍，使硬件成本降低 10 余倍；

• 一行代码实现自动搜索最佳并行策略，显著降低分布式训练上手门槛，原生支持 Hugging Face，Timm 等热门 AI 模型库。

然而，高昂的成本也一定程度上限制了 AIGC 产业的大规模普及。以 Stable Diffusion 为例，其背后的 Stability AI 维护超过 4000 个英伟达 A100 的 GPU 集群，并已为此支出超过 5000 万美元的运营成本 。 面对快速迭代的模型、算法和下游任务，如何降低应用成本成为 AIGC 真正走向落地的核心问题。

Stable Diffusion 2.0 基于简单易用的 PyTorch Lightning 框架搭建。 作为 PyTorch Lightning 的官方大模型解决方案 ，Colossal-AI 在第一时间跟进，开源了更为高效，硬件门槛更低的全套训练 / 微调 / 推理方案:

• 可减少训练显存消耗 5.6 倍，硬件成本最高降低 46 倍；
• 支持 DreamBooth 单 GPU 快速个性化微调；

• 推理显存消耗降低 2.5 倍。

Colossal-AI 通过一系列显存优化技术和支持 Stable Diffusion 2.0，使 Stable Diffusion 平均在每个 GPU 上使用大 batch size 16 训练的 显存需求从 64.5GB 降低 5.6 倍至 11.6GB ，还可扩展至单 GPU 或多 GPU 并行。相比使用最先进的 A100 80GB，目前仅需 3060 等消费级显卡即可满足需求， 硬件成本最高直降 46 倍 。更多用户可以在消费级 GPU 上，低成本地开展 Stable Diffusion 的相关研究与应用落地。

1.1.1 背后显存优化

Flash Attention： 早在 Stable Diffusion 1.0 版本，Colossal-AI 就率先引入了 Flash Attention 技术，成功将 attention 的速度提升 104%，将端到端训练的峰值显存减少 23%。Flash Attention 是针对长序列 attention 的加速版本，使用 Flatten 来减少 GPU 高带宽内存（HBM）之间的内存读 / 写次数，Flash attention 同时针对块状稀疏的 attention，设计了一个近似的注意力算法，比任何现有的近似 attention 方法都要快。在 Stable Diffusion 1.0 版本，整个 Diffusion Model 只有少量 attention 层，Flash attention 还没有体现出其性能优势。在 Stable Diffusion 2.0 中，由于将大量卷积层替换为 attention 层，进一步发挥了 Flash Attention 的显存优化潜力。

ZeRO + Gemini： Colossal-AI 支持使用零冗余优化器 (ZeRO) 的方法来消除内存冗余，与经典的数据并行性策略相比，可极大提高内存使用效率，同时不牺牲计算粒度和通信效率。此外，Colossal-AI 还引入了 Chunk 机制进一步提升 ZeRO 的性能。运算顺序上连续的一组参数存入一个 Chunk 中（Chunk 即一段连续的内存空间），每个 Chunk 的大小相同。Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU 之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。

1.1.2 一行代码快速上手

作为 PyTorch Lightning 的官方合作伙伴， 仅需一行代码即可调用 ColossalAI 的上述显存优化 。

from lightning.pytorch import trainer, LightningModulefrom lightning.pytorch.strategies import ColossalAIStrategyMystrategy = ColossalAIStrategy(use_chunk=True, enable_distributed_storage=True, placement_policy=auto)trainer = Trainer(accelerator="gpu", devices=4, precision=16, strategy=Mystrategy)trainer.fit(model)

DreamBooth 是一种用于个性化文本到图像模型（例如 Stable Diffusion）的方法， 只需要几张（3~5 张）期望主题的图像 ，即可生成指定主题的一系列图像。用户只需直接运行文件 train_dreambooth_colossalai.py，即可在该微调任务上充分发挥 Colossal-AI 的显存优化，个性化快速微调自己的图文模型，极大地降低了使用门槛。

由于模型推理对数值精度不敏感，这为实现低精度的低成本推理提供了可能。对于 Stable Diffusion 2.0 模型，可以通过 添加一行代码 ，支持模型的 Int8 量化推理， 显存消耗降低 2.5 倍，仅需 3.1GB ，且不造成显著性能损失。

model = replace_module(model)

随着模型规模的不断增大，模型推理的显存及内存占用也变成了不可忽视的关键因素。以 Hugging Face 发布的 1750 亿参数开源模型 BLOOM 为例，如果直接使用常见的 FP32/FP16 进行推理，在单节点 8 张 GPU 使用模型并行， 每张 GPU 需要消耗至少 87.5GB/43.8GB 的显存 。如此大的显存占用，即使是最先进的 8 卡 A100（80GB/40GB）服务器，也无法直接部署推理服务，而多节点推理又会带来沉重的额外成本和通信开销。

Colossal-AI 实现了高效的 Int8 量化和模型并行推理，可以 将 1750 亿参数的 BLOOM 等大模型的推理服务，部署到 3090/4090 等消费级显卡的 8 卡服务器 ，同时不产生显著的 CPU 内存占用提升及性能损耗。相比原有的 A100 方案， 可降低硬件部署成本十余倍 。

在 Colossal-AI 中，通过对模型进行了 Int8 量化，可将模型总体显存占用从 352.3GB（FP16）降低到 185.6GB， 同时使用 Colossal-AI 的模型并行技术， 将每张显卡的占用减少到了 23.2GB 。在模型并行中，为了不增加 CPU 内存占用，Colossal-AI 在主进程中进行模型的量化和切分工作，其余的每个进程中分别使用 lazy_init 获得几乎不占显存和内存的 meta model，再通过 gloo backend 在进程之间传递模型参数。通过上述方案，在没有分段加载模型参数的情况下，便可以使得 CPU 内存占用峰值，达到理论上的较优水平。相较于将模型按层切分的 “类流水线” 分布方式，模型并行可以提高非密集请求下的显存使用效率。

Colossal-AI 提供了业界急需的一套高效易用 自动并行系统 。相比现有其他手动配置复杂并行策略和修改模型的解决方案， Colossal-AI 仅需增加一行代码 ，提供 cluster 信息以及单机训练模型即可获得分布式训练能力，并且 原生支持包括 Hugging Face，Timm 等热门 AI 模型库 。

# wrap the model using auto_enginemodel, optimizer = auto_engine(model, optimizer, cluster_info)# normal training loop...

因此，Colossal-AI 可以 极大地降低 AI 开发者的使用分布式技术训练和微调大模型门槛 。同时，自动并行系统可以从更细粒度搜索并行策略， 找到更加高效的并行方案 。

Colossal-AI 是 首个基于 PyTorch 框架使用静态图分析的自动并行系统 。PyTorch 作为一个动态图框架，获取其静态的执行计划是机器学习系统领域被长期研究的问题。Colossal-AI 使用基于 torch.FX Tracer 的 ColoTracer，在 tracing 过程中推导并记录了每个 tensor 的元信息，例如 tensor shape，dims，dtype 等，可以为后续的自动并行策略搜索提供帮助。因此 Colossal-AI 具有更好的模型泛化能力，而不是依靠模型名或手动修改来适配并行策略。

1. Colossal-AI 的 device mesh 可以 profiling 到集群性能指标，对不同的通信算子进行耗时估算。
2. Colossal-AI 的 shape consistency 会贪心地搜索 sharding spec 间的转换方式，而不是朴素地逐 dimension 进行转换，这样能找到更高效的转换路径，进而使得 sharding spec 间的转换通信开销更小。
3. 加入了 all_to_all 操作，使得 Colossal-AI 的扩展性更强，这在大规模集群上进行训练时，可以展现出很大的优势。