明敏发自凹非寺
量子位公众号 QbitAI
从 AI 画画到 NLP 大模型，AIGC 的落地成本被一次性打下来了！
话不多说，直接看结果：
- Stable Diffusion 2.0 训练/微调/推理，显存消耗最多可节省 5.6 倍，使硬件成本直降至1/46，一行代码即可启用；
- 1750 亿参数大模型 BLOOM 单机推理，显存消耗节省 4 倍，硬件成本直降至十几分之一
- 一行代码实现自动搜索最佳并行策略，显著降低分布式训练上手门槛，原生支持 Hugging Face、Timm 等热门 AI 模型库。
要知道，在 AIGC 爆火的另一面，居高不下的成本困扰着整个行业。
上周，首批 AI 画画公司之一 StockAI 被迫宣布关闭平台。原因无他，创始人表示：公司驱动成本太高了，目前的收入难以为继。
即便如 ChatGPT 身后有 OpenAI 和微软撑腰，也在平台开放几周后发出公告，开始限制用户每日使用次数。
言下之意无非四个字：烧不起了。
总而言之，降低 AI 大模型落地成本，是目前行业内亟需解决的问题。
与此同时，开源 AI 大模型解决方案 Colossal-AI 在过去一年迅速蹿红，GitHub 上已经收获 7k+Star。
如上降本方案，便出自它之手。
具体是怎么实现的？往下看，
开源地址：
https://github.com/hpcaitech/ColossalAI
Stable Diffusion2.0 低成本训练/微调/推理
相较于 1.0 版本，Stable Diffusion 2.0 不仅提升了模型生成图像像素，还引入了 Depth2img 模型、text-guided 修复模型等，功能更加完善。
这波上新其实让用户们既惊喜又措手不及。
（毕竟 1.0 都还没玩明白呢）
但话说回来，还是老问题，AIGC 模型落地的成本高啊。
以 Stable Diffusion 为例，其背后的 Stability AI 维护超过 4000 个英伟达 A100 的 GPU 集群，并已为此支出超过 5000 万美元的运营成本。
面对快速迭代的模型、算法和下游任务，如何降低应用成本成为 AIGC 真正走向落地的核心问题。
Stable Diffusion 2.0 基于简单易用的 PyTorch Lightning 框架搭建。
作为 PyTorch Lightning 的官方大模型解决方案，Colossal-AI 第一时间进行跟进。
具体内容有以下几点：
- 显存消耗可节省 5.6 倍，硬件成本最多降至1/46
- 支持 DreamBooth 单 GPU 快速个性化微调
- 推理显存消耗节省 2.5 倍
而且该方案也将于近期合并进入 Hugging Face，进一步方便用户使用。
训练
为了加快训练速度，降低训练成本，使用更大的 batch size 已成为被广泛使用的有效手段。但 GPU 有限的显存容量，严重限制了 batch size 大小，推高了训练硬件门槛。
通过一系列显存优化技术，Colossal-AI 使得 Stable Diffusion 平均在每个 GPU 上使用大 batch size 16 训练的显存需求，从 64.5GB 降低到了 11.6GB、节省 5.6 倍，还可扩展至单 GPU 或多 GPU 并行。
相比使用最先进的 A100 80GB，目前仅需3060 等消费级显卡即可满足需求，硬件成本最高直降至1/46。
由此可以让更多用户在消费级 GPU 上，低成本地开展 Stable Diffusion 的相关研究与应用落地。
背后显存优化
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 内存墙的限制，进一步扩展了可训练模型规模。
一行代码快速上手
作为 PyTorch Lightning 的官方合作伙伴，仅需一行代码即可调用 Colossal-AI 的上述显存优化。
fromlightning.pytorchimporttrainer,LightningModule
fromlightning.pytorch.strategiesimportColossalAIStrategy
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)
DreamBooth 微调
在推出 Stable Diffusion 2.0 加速方案的同时，Colossal-AI 还“顺手”发布了 DreamBooth 模型的微调方案。
这是谷歌在今年 8 月发布的模型。它只需3-5 张图片，再加上文字表述，就能让指定物体迁移到其他场景或风格中去。
和 Dall-E 2、Imagen 等最大的不同是，DreamBooth 能对选定对象忠实还原。
方案中，用户只需直接运行文件 train_dreambooth_colossalai.py，即可在该微调任务上充分发挥 Colossal-AI 的显存优化，个性化快速微调自己的图文模型，极大降低使用门槛。
推理
由于模型推理对数值精度不敏感，这为实现低精度的低成本推理提供了可能。
对于 Stable Diffusion 2.0 模型，可以通过添加一行代码，支持模型的 Int8 量化推理，显存消耗节省 2.5 倍，仅需 3.1GB，且不造成显著性能损失。
model=replace_module (model)
用 RTX3090 推理 1750 亿 BLOOM 模型
AI 画画爆火的另一边，NLP 大模型的趋势也还在延续。
今年 7 月，Hugging Face 发布了 1750 亿参数开源模型 BLOOM，它动用 384 块 A100 炼成。
如果直接使用常见的 FP32/FP16 进行推理，在单节点 8 张 GPU 使用模型并行，每张 GPU 需要消耗至少 87.5GB/43.8GB 的显存。
如此大的显存占用，即使是最先进的 8 卡 A100（80GB/40GB）服务器，也无法直接部署推理服务，而多节点推理又会带来沉重的额外成本和通信开销。
基于这一现状，Colossal-AI 实现了高效的 Int8 量化和模型并行推理，可以将 1750 亿参数的 BLOOM 等大模型的推理服务，部署到 3090/4090 等消费级显卡的 8 卡服务器，同时不产生显著的 CPU 内存占用提升及性能损耗。
相比原有的 A100 方案，可将硬件部署成本降低到原有的十几分之一。
通过对模型进行 Int8 量化，Colossal-AI 可将模型总体显存占用从 352.3GB（FP16）降低到 185.6GB， 同时使用 Colossal-AI 的模型并行技术，将每张显卡的占用减少到了 23.2GB。
在模型并行中，为了不增加 CPU 内存占用，Colossal-AI 在主进程中进行模型的量化和切分工作，其余的每个进程中分别使用 lazy_init 获得几乎不占显存和内存的 meta model，再通过 gloo backend 在进程之间传递模型参数。
通过上述方案，在没有分段加载模型参数的情况下，便可以使得 CPU 内存占用峰值，达到理论上的较优水平。相较于将模型按层切分的“类流水线”分布方式，模型并行可以提高非密集请求下的显存使用效率。
一行代码自动并行
大模型的分布式混合部署是一个非常复杂的问题。
目前常见的分布式大模型训练方案，都依赖于用户人工反复尝试以及系统专家基于经验进行配置部署。
然而，这对于绝大多数 AI 开发者来说很不友好，因为大家都不希望把过多时间精力花费在研究分布式系统和试错上。
由此，Colossal-AI 的高效易用自动并行系统，可以说是解大家燃眉之急了。
仅需增加一行代码，它就能提供 cluster 信息以及单机训练模型即可获得分布式训练能力，并且原生支持包括 Hugging Face，Timm 等热门 AI 模型库。
#wrapthemodelusingauto_engine
model,optimizer=auto_engine (model,optimizer,cluster_info)
#normaltrainingloop
...
因此，Colossal-AI 可以极大地降低 AI 开发者的使用分布式技术训练和微调大模型门槛。同时，自动并行系统可以从更细粒度搜索并行策略，找到更加高效的并行方案。
Graph Tracing
Colossal-AI 是首个基于 PyTorch 框架使用静态图分析的自动并行系统。
PyTorch 作为一个动态图框架，获取其静态的执行计划是机器学习系统领域被长期研究的问题。
Colossal-AI 使用基于 torch.FX Tracer 的 ColoTracer，在 tracing 过程中推导并记录了每个 tensor 的元信息，例如 tensor shape，dims，dtype 等，可以为后续的自动并行策略搜索提供帮助。
因此 Colossal-AI 具有更好的模型泛化能力，而不是依靠模型名或手动修改来适配并行策略。
细粒度分布式训练策略搜索
Colossal-AI 会在满足内存预算的限制下，以最快运行时间为目标，为每个 op 进行策略搜索，最终得到真实训练时的策略，包括每个 tensor 的切分策略，不同计算节点间需要插入的通信算子类型，是否要进行算子替换等。
现有系统中的张量并行、数据并行，NVIDIA 在 Megatron-LM 等并行系统中使用的 column 切分和 row 切分并行等混合并行，都是自动并行可以搜索到的策略的子集。
除了这些可以手动指定的并行方式外，Colossal-AI 的自动并行系统有能力为每个 op 指定独特的并行方式，因此有可能找到比依赖专家经验和试错配置的手动切分更好的并行策略。
分布式 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 的扩展性更强，这在大规模集群上进行训练时，可以展现出很大的优势
与 activation checkpoint 结合
作为大模型训练中必不可少的显存压缩技术，Colossal-AI 也提供了对于 activation checkpoint 的自动搜索功能。
相比于大部分将最大显存压缩作为目标的技术方案，Colossal-AI 的搜索目标是在显存预算以内，找到最快的 activation checkpoint 方案。
同时，为了避免将 activation checkpoint 的搜索一起建模到 SPMD solver 中导致搜索时间爆炸，Colossal-AI 做了2-stage search 的设计，因此可以在合理的时间内搜索到有效可行的分布式训练方案。
关于 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
参考链接：
