硬件预算最高直降46倍!低成本AIGC和千亿大模型,一行代码自动并行
机器之心发布
机器之心编辑部
最近，AI 大模型连续火爆出圈，人工智能生成模型（AIGC）的热度尚未褪去，聊天机器人 ChatGPT 便引发全网热议，两周吸引百万用户。还有卷趴一半程序员的 AlphaCode，生成全新蛋白质的 ESM2 等，不断探索 AI 大模型落地的新领域。面对大模型带来的技术革命，连谷歌都拉响 “红色警报”，担心 ChatGPT 砸掉自己搜索引擎的饭碗。
作为当下最火热的开源 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 模型库。
开源地址：
https://github.com/hpcaitech/ColossalAI
1. Stable Diffusion 2.0 低成本训练 / 微调 / 推理
AIGC（人工智能生成内容）目前是 AI 领域最热门的话题之一，还被 Science 评为 2022 年度十大科学突破。短短数月，爆红 AI 圈的 Stable Diffusion 已升级到 2.0 版本，正如部分用户所言：“我从未见过任何技术发展如此之快。V1 还没整透彻，V2 就来了”。
然而，高昂的成本也一定程度上限制了 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 倍。
该方案也将于近期合并进入当下最火热的 AI 模型社区 Hugging Face，进一步方便用户使用。
1.1 训练
为了加快训练速度，降低训练成本，使用更大的 batch size 已成为被广泛使用的有效手段。但 GPU 有限的显存容量，严重限制了 batch size 大小，推高了训练硬件门槛。
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 之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。
而 Colossal-AI 的异构内存空间管理器 Gemini 支持将优化器状态从 GPU 卸载到 CPU，以节省 GPU 内存占用。可以同时利用 GPU 内存、CPU 内存（由 CPU DRAM 或 NVMe SSD 内存组成）来突破单 GPU 内存墙的限制，进一步扩展了可训练模型规模。
1.1.2 一行代码快速上手
作为 PyTorch Lightning 的官方合作伙伴，仅需一行代码即可调用 ColossalAI 的上述显存优化。
from lightning.pytorch import trainer, LightningModule
from lightning.pytorch.strategies import ColossalAIStrategy
Mystrategy = ColossalAIStrategy(use_chunk=True,
enable_distributed_storage=True, placement_policy=auto)
trainer = Trainer(accelerator="gpu", devices=4, precision=16,
strategy=Mystrategy)
trainer.fit(model)
1.2 DreamBooth 微调
DreamBooth 是一种用于个性化文本到图像模型（例如 Stable Diffusion）的方法，只需要几张（3~5 张）期望主题的图像，即可生成指定主题的一系列图像。用户只需直接运行文件 train_dreambooth_colossalai.py，即可在该微调任务上充分发挥 Colossal-AI 的显存优化，个性化快速微调自己的图文模型，极大地降低了使用门槛。
1.3 推理
由于模型推理对数值精度不敏感，这为实现低精度的低成本推理提供了可能。对于 Stable Diffusion 2.0 模型，可以通过添加一行代码，支持模型的 Int8 量化推理，显存消耗降低 2.5 倍，仅需 3.1GB，且不造成显著性能损失。
model = replace_module(model)
2. 1750 亿 BLOOM 模型低成本推理
随着模型规模的不断增大，模型推理的显存及内存占用也变成了不可忽视的关键因素。以 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 内存占用峰值，达到理论上的较优水平。相较于将模型按层切分的 “类流水线” 分布方式，模型并行可以提高非密集请求下的显存使用效率。
3. 一行代码自动并行
大模型的分布式混合部署是一个非常复杂的问题，目前常见的分布式大模型训练方案，都依赖用户人工反复尝试以及系统专家的经验来进行配置部署。然而这对绝大多数 AI 开发者来说十分不友好，因为不希望将时间精力花费在研究分布式系统和试错上。
Colossal-AI 提供了业界急需的一套高效易用自动并行系统。相比现有其他手动配置复杂并行策略和修改模型的解决方案，Colossal-AI 仅需增加一行代码，提供 cluster 信息以及单机训练模型即可获得分布式训练能力，并且原生支持包括 Hugging Face，Timm 等热门 AI 模型库。
# wrap the model using auto_engine
model, optimizer = auto_engine(model, optimizer, cluster_info)
# normal training loop
因此，Colossal-AI 可以极大地降低 AI 开发者的使用分布式技术训练和微调大模型门槛。同时，自动并行系统可以从更细粒度搜索并行策略，找到更加高效的并行方案。
3.1 Graph Tracing
Colossal-AI 是首个基于 PyTorch 框架使用静态图分析的自动并行系统。PyTorch 作为一个动态图框架，获取其静态的执行计划是机器学习系统领域被长期研究的问题。Colossal-AI 使用基于 torch.FX Tracer 的 ColoTracer，在 tracing 过程中推导并记录了每个 tensor 的元信息，例如 tensor shape，dims，dtype 等，可以为后续的自动并行策略搜索提供帮助。因此 Colossal-AI 具有更好的模型泛化能力，而不是依靠模型名或手动修改来适配并行策略。
3.2 细粒度分布式训练策略搜索
Colossal-AI 会在满足内存预算的限制下，以最快运行时间为目标，为每个 op 进行策略搜索，最终得到真实训练时的策略，包括每个 tensor 的切分策略，不同计算节点间需要插入的通信算子类型，是否要进行算子替换等。现有系统中的张量并行，数据并行，NVIDIA 在 Megatron-LM 等并行系统中使用的 column 切分和 row 切分并行等混合并行，都是自动并行可以搜索到的策略的子集。除了这些可以手动指定的并行方式外，Colossal-AI 的自动并行系统有能力为每个 op 指定独特的并行方式，因此有可能找到比依赖专家经验和试错配置的手动切分更好的并行策略。
3.3 分布式 tensor 与 shape consistency 系统
与 PyTorch 最新发布的 DTensor 类似，Colossal-AI 也使用了 device mesh 对集群进行了抽象管理。具体来说，Colossal-AI 使用 sharding spec 对 tensor 的分布式存储状态进行标注，使用 shape consistency manager 自动地对同一 tensor 在不同 sharding spec 间进行转换。这让 Colossal-AI 的通用性和易用性极大地提升，借助 shape consistency manager 可以没有负担地切分 tensor，而不用担心上游 op 的 output 与下游的 input 在集群中的存储方式不同。
相较于 PyTorch DTensor，Colossal-AI 有以下 3 个优势：
- Colossal-AI 的 device mesh 可以 profiling 到集群性能指标，对不同的通信算子进行耗时估算。
- Colossal-AI 的 shape consistency 会贪心地搜索 sharding spec 间的转换方式，而不是朴素地逐 dimension 进行转换，这样能找到更高效的转换路径，进而使得 sharding spec 间的转换通信开销更小。
- 加入了 all_to_all 操作，使得 Colossal-AI 的扩展性更强，这在大规模集群上进行训练时，可以展现出很大的优势。
3.4 与 activation checkpoint 结合
作为大模型训练中必不可少的显存压缩技术，Colossal-AI 也提供了对于 activation checkpoint 的自动搜索功能。相比于大部分将最大显存压缩作为目标的技术方案，Colossal-AI 的搜索目标是在显存预算以内，找到最快的 activation checkpoint 方案。同时，为了避免将 activation checkpoint 的搜索一起建模到 SPMD solver 中导致搜索时间爆炸，Colossal-AI 做了 2-stage search 的设计，因此可以在合理的时间内搜索到有效可行的分布式训练方案。
4. 关于 Colossal-AI
面向大模型时代的通用深度学习系统 Colossal-AI，可实现高效快速部署 AI 大模型训练和推理，降低 AI 大模型应用成本。自开源以来，Colossal-AI 已经多次在 GitHub 热榜位列世界第一，获得 GitHub Star 超七千颗，并成功入选 SC、AAAI、PPoPP 等国际 AI 与 HPC 顶级会议的官方教程。
Colossal-AI 相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用，广受好评。例如，近期火爆的 ChatGPT 尚未开源，且不具备联网功能。Colossal-AI 已成功帮助某世界 500 强企业，开发具备在线搜索引擎能力增强的聊天机器人模型。
开源地址：
https://github.com/hpcaitech/ColossalAI
参考链接：
https://www.hpc-ai.tech/blog/colossal-ai-0-2-0